import argparse
import json

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import SegformerImageProcessor, SwinConfig, UperNetConfig, UperNetForSemanticSegmentation





def      lowerCAmelCase_       ( __lowerCAmelCase  )->     Optional[Any]:






      '''simple docstring'''





      UpperCAmelCase			:       Dict			     =3_84
      UpperCAmelCase			:       Tuple			     =7
      if "tiny" in model_name:
            UpperCAmelCase			:       List[str]			     =96
            UpperCAmelCase			:       Any			     =(2, 2, 6, 2)
            UpperCAmelCase			:       List[str]			     =(3, 6, 12, 24)
      elif "small" in model_name:
            UpperCAmelCase			:       List[Any]			     =96
            UpperCAmelCase			:       str			     =(2, 2, 18, 2)
            UpperCAmelCase			:       List[str]			     =(3, 6, 12, 24)
      elif "base" in model_name:
            UpperCAmelCase			:       Dict			     =1_28
            UpperCAmelCase			:       Tuple			     =(2, 2, 18, 2)
            UpperCAmelCase			:       str			     =(4, 8, 16, 32)
            UpperCAmelCase			:       int			     =12
            UpperCAmelCase			:       int			     =5_12
      elif "large" in model_name:
            UpperCAmelCase			:       Optional[Any]			     =1_92
            UpperCAmelCase			:       Any			     =(2, 2, 18, 2)
            UpperCAmelCase			:       List[str]			     =(6, 12, 24, 48)
            UpperCAmelCase			:       Tuple			     =12
            UpperCAmelCase			:       int			     =7_68

      # set label information
      UpperCAmelCase			:       List[Any]			     =1_50
      UpperCAmelCase			:       Dict			     ='''huggingface/label-files'''
      UpperCAmelCase			:       List[str]			     ='''ade20k-id2label.json'''
      UpperCAmelCase			:       List[str]			     =json.load(open(hf_hub_download(__lowerCAmelCase						,    __lowerCAmelCase						,    repo_type='''dataset'''  )						,    '''r'''  )  )
      UpperCAmelCase			:       Tuple			     ={int(__lowerCAmelCase  ): v for k, v in idalabel.items()}
      UpperCAmelCase			:       Any			     ={v: k for k, v in idalabel.items()}

      UpperCAmelCase			:       str			     =SwinConfig(
          embed_dim=__lowerCAmelCase						,    depths=__lowerCAmelCase						,    num_heads=__lowerCAmelCase						,    window_size=__lowerCAmelCase						,    out_features=['''stage1''', '''stage2''', '''stage3''', '''stage4''']						,    )
      UpperCAmelCase			:       Any			     =UperNetConfig(
          backbone_config=__lowerCAmelCase						,    auxiliary_in_channels=__lowerCAmelCase						,    num_labels=__lowerCAmelCase						,    idalabel=__lowerCAmelCase						,    labelaid=__lowerCAmelCase						,    )

      return config





def      lowerCAmelCase_       ( __lowerCAmelCase  )->     List[str]:






      '''simple docstring'''





      UpperCAmelCase			:       Tuple			     =[]

      # fmt: off
      # stem
      rename_keys.append(('''backbone.patch_embed.projection.weight''', '''backbone.embeddings.patch_embeddings.projection.weight''')  )
      rename_keys.append(('''backbone.patch_embed.projection.bias''', '''backbone.embeddings.patch_embeddings.projection.bias''')  )
      rename_keys.append(('''backbone.patch_embed.norm.weight''', '''backbone.embeddings.norm.weight''')  )
      rename_keys.append(('''backbone.patch_embed.norm.bias''', '''backbone.embeddings.norm.bias''')  )
      # stages
      for i in range(len(config.backbone_config.depths  )  ):
            for j in range(config.backbone_config.depths[i]  ):
                  rename_keys.append((f'''backbone.stages.{i}.blocks.{j}.norm1.weight''', f'''backbone.encoder.layers.{i}.blocks.{j}.layernorm_before.weight''')  )
                  rename_keys.append((f'''backbone.stages.{i}.blocks.{j}.norm1.bias''', f'''backbone.encoder.layers.{i}.blocks.{j}.layernorm_before.bias''')  )
                  rename_keys.append((f'''backbone.stages.{i}.blocks.{j}.attn.w_msa.relative_position_bias_table''', f'''backbone.encoder.layers.{i}.blocks.{j}.attention.self.relative_position_bias_table''')  )
                  rename_keys.append((f'''backbone.stages.{i}.blocks.{j}.attn.w_msa.relative_position_index''', f'''backbone.encoder.layers.{i}.blocks.{j}.attention.self.relative_position_index''')  )
                  rename_keys.append((f'''backbone.stages.{i}.blocks.{j}.attn.w_msa.proj.weight''', f'''backbone.encoder.layers.{i}.blocks.{j}.attention.output.dense.weight''')  )
                  rename_keys.append((f'''backbone.stages.{i}.blocks.{j}.attn.w_msa.proj.bias''', f'''backbone.encoder.layers.{i}.blocks.{j}.attention.output.dense.bias''')  )
                  rename_keys.append((f'''backbone.stages.{i}.blocks.{j}.norm2.weight''', f'''backbone.encoder.layers.{i}.blocks.{j}.layernorm_after.weight''')  )
                  rename_keys.append((f'''backbone.stages.{i}.blocks.{j}.norm2.bias''', f'''backbone.encoder.layers.{i}.blocks.{j}.layernorm_after.bias''')  )
                  rename_keys.append((f'''backbone.stages.{i}.blocks.{j}.ffn.layers.0.0.weight''', f'''backbone.encoder.layers.{i}.blocks.{j}.intermediate.dense.weight''')  )
                  rename_keys.append((f'''backbone.stages.{i}.blocks.{j}.ffn.layers.0.0.bias''', f'''backbone.encoder.layers.{i}.blocks.{j}.intermediate.dense.bias''')  )
                  rename_keys.append((f'''backbone.stages.{i}.blocks.{j}.ffn.layers.1.weight''', f'''backbone.encoder.layers.{i}.blocks.{j}.output.dense.weight''')  )
                  rename_keys.append((f'''backbone.stages.{i}.blocks.{j}.ffn.layers.1.bias''', f'''backbone.encoder.layers.{i}.blocks.{j}.output.dense.bias''')  )

            if i < 3:
                  rename_keys.append((f'''backbone.stages.{i}.downsample.reduction.weight''', f'''backbone.encoder.layers.{i}.downsample.reduction.weight''')  )
                  rename_keys.append((f'''backbone.stages.{i}.downsample.norm.weight''', f'''backbone.encoder.layers.{i}.downsample.norm.weight''')  )
                  rename_keys.append((f'''backbone.stages.{i}.downsample.norm.bias''', f'''backbone.encoder.layers.{i}.downsample.norm.bias''')  )
            rename_keys.append((f'''backbone.norm{i}.weight''', f'''backbone.hidden_states_norms.stage{i+1}.weight''')  )
            rename_keys.append((f'''backbone.norm{i}.bias''', f'''backbone.hidden_states_norms.stage{i+1}.bias''')  )

      # decode head
      rename_keys.extend(
          [
              ('''decode_head.conv_seg.weight''', '''decode_head.classifier.weight'''),
              ('''decode_head.conv_seg.bias''', '''decode_head.classifier.bias'''),
              ('''auxiliary_head.conv_seg.weight''', '''auxiliary_head.classifier.weight'''),
              ('''auxiliary_head.conv_seg.bias''', '''auxiliary_head.classifier.bias'''),
          ]  )
      # fmt: on

      return rename_keys





def      lowerCAmelCase_       ( __lowerCAmelCase						,    __lowerCAmelCase						,    __lowerCAmelCase  )->     Optional[int]:






      '''simple docstring'''





      UpperCAmelCase			:       Optional[Any]			     =dct.pop(__lowerCAmelCase  )
      UpperCAmelCase			:       Optional[Any]			     =val





def      lowerCAmelCase_       ( __lowerCAmelCase						,    __lowerCAmelCase  )->     Any:






      '''simple docstring'''





      UpperCAmelCase			:       Union[str, Any]			     =[int(backbone_config.embed_dim * 2**i  ) for i in range(len(backbone_config.depths  )  )]
      for i in range(len(backbone_config.depths  )  ):
            UpperCAmelCase			:       Tuple			     =num_features[i]
            for j in range(backbone_config.depths[i]  ):
                  # fmt: off
                  # read in weights + bias of input projection layer (in original implementation, this is a single matrix + bias)
                  UpperCAmelCase			:       int			     =state_dict.pop(f'''backbone.stages.{i}.blocks.{j}.attn.w_msa.qkv.weight'''  )
                  UpperCAmelCase			:       Union[str, Any]			     =state_dict.pop(f'''backbone.stages.{i}.blocks.{j}.attn.w_msa.qkv.bias'''  )
                  # next, add query, keys and values (in that order) to the state dict
                  UpperCAmelCase			:       Tuple			     =in_proj_weight[:dim, :]
                  UpperCAmelCase			:       List[Any]			     =in_proj_bias[: dim]
                  UpperCAmelCase			:       Union[str, Any]			     =in_proj_weight[
                      dim : dim * 2, :
                  ]
                  UpperCAmelCase			:       int			     =in_proj_bias[
                      dim : dim * 2
                  ]
                  UpperCAmelCase			:       Optional[int]			     =in_proj_weight[
                      -dim :, :
                  ]
                  UpperCAmelCase			:       int			     =in_proj_bias[-dim :]
                  # fmt: on





def      lowerCAmelCase_       ( __lowerCAmelCase  )->     List[str]:






      '''simple docstring'''





      UpperCAmelCase     ,  UpperCAmelCase			:       Optional[int]			     =x.shape
      UpperCAmelCase			:       str			     =x.reshape(__lowerCAmelCase						,    4						,    in_channel // 4  )
      UpperCAmelCase			:       Tuple			     =x[:, [0, 2, 1, 3], :].transpose(1						,    2  ).reshape(__lowerCAmelCase						,    __lowerCAmelCase  )
      return x





def      lowerCAmelCase_       ( __lowerCAmelCase  )->     Any:






      '''simple docstring'''





      UpperCAmelCase     ,  UpperCAmelCase			:       List[str]			     =x.shape
      UpperCAmelCase			:       Dict			     =x.reshape(__lowerCAmelCase						,    in_channel // 4						,    4  )
      UpperCAmelCase			:       Union[str, Any]			     =x[:, :, [0, 2, 1, 3]].transpose(1						,    2  ).reshape(__lowerCAmelCase						,    __lowerCAmelCase  )

      return x





def      lowerCAmelCase_       ( __lowerCAmelCase  )->     Any:






      '''simple docstring'''





      UpperCAmelCase			:       Union[str, Any]			     =x.shape[0]
      UpperCAmelCase			:       Optional[Any]			     =x.reshape(4						,    in_channel // 4  )
      UpperCAmelCase			:       Dict			     =x[[0, 2, 1, 3], :].transpose(0						,    1  ).reshape(__lowerCAmelCase  )
      return x





def      lowerCAmelCase_       ( __lowerCAmelCase  )->     List[str]:






      '''simple docstring'''





      UpperCAmelCase			:       Tuple			     =x.shape[0]
      UpperCAmelCase			:       int			     =x.reshape(in_channel // 4						,    4  )
      UpperCAmelCase			:       Dict			     =x[:, [0, 2, 1, 3]].transpose(0						,    1  ).reshape(__lowerCAmelCase  )
      return x





def      lowerCAmelCase_       ( __lowerCAmelCase						,    __lowerCAmelCase						,    __lowerCAmelCase  )->     Any:






      '''simple docstring'''





      UpperCAmelCase			:       List[Any]			     ={
          '''upernet-swin-tiny''': '''https://download.openmmlab.com/mmsegmentation/v0.5/swin/upernet_swin_tiny_patch4_window7_512x512_160k_ade20k_pretrain_224x224_1K/upernet_swin_tiny_patch4_window7_512x512_160k_ade20k_pretrain_224x224_1K_20210531_112542-e380ad3e.pth''',
          '''upernet-swin-small''': '''https://download.openmmlab.com/mmsegmentation/v0.5/swin/upernet_swin_small_patch4_window7_512x512_160k_ade20k_pretrain_224x224_1K/upernet_swin_small_patch4_window7_512x512_160k_ade20k_pretrain_224x224_1K_20210526_192015-ee2fff1c.pth''',
          '''upernet-swin-base''': '''https://download.openmmlab.com/mmsegmentation/v0.5/swin/upernet_swin_base_patch4_window12_512x512_160k_ade20k_pretrain_384x384_22K/upernet_swin_base_patch4_window12_512x512_160k_ade20k_pretrain_384x384_22K_20210531_125459-429057bf.pth''',
          '''upernet-swin-large''': '''https://download.openmmlab.com/mmsegmentation/v0.5/swin/upernet_swin_large_patch4_window12_512x512_pretrain_384x384_22K_160k_ade20k/upernet_swin_large_patch4_window12_512x512_pretrain_384x384_22K_160k_ade20k_20220318_091743-9ba68901.pth''',
      }
      UpperCAmelCase			:       str			     =model_name_to_url[model_name]
      UpperCAmelCase			:       Any			     =torch.hub.load_state_dict_from_url(__lowerCAmelCase						,    map_location='''cpu'''						,    file_name=__lowerCAmelCase  )[
          '''state_dict'''
      ]

      for name, param in state_dict.items():
            print(__lowerCAmelCase						,    param.shape  )

      UpperCAmelCase			:       Optional[Any]			     =get_upernet_config(__lowerCAmelCase  )
      UpperCAmelCase			:       Optional[int]			     =UperNetForSemanticSegmentation(__lowerCAmelCase  )
      model.eval()

      # replace "bn" => "batch_norm"
      for key in state_dict.copy().keys():
            UpperCAmelCase			:       int			     =state_dict.pop(__lowerCAmelCase  )
            if "bn" in key:
                  UpperCAmelCase			:       str			     =key.replace('''bn'''						,    '''batch_norm'''  )
            UpperCAmelCase			:       Any			     =val

      # rename keys
      UpperCAmelCase			:       Union[str, Any]			     =create_rename_keys(__lowerCAmelCase  )
      for src, dest in rename_keys:
            rename_key(__lowerCAmelCase						,    __lowerCAmelCase						,    __lowerCAmelCase  )
      read_in_q_k_v(__lowerCAmelCase						,    config.backbone_config  )

      # fix downsample parameters
      for key, value in state_dict.items():
            if "downsample" in key:
                  if "reduction" in key:
                        UpperCAmelCase			:       str			     =reverse_correct_unfold_reduction_order(__lowerCAmelCase  )
                  if "norm" in key:
                        UpperCAmelCase			:       int			     =reverse_correct_unfold_norm_order(__lowerCAmelCase  )

      model.load_state_dict(__lowerCAmelCase  )

      # verify on image
      UpperCAmelCase			:       List[Any]			     ='''https://huggingface.co/datasets/hf-internal-testing/fixtures_ade20k/resolve/main/ADE_val_00000001.jpg'''
      UpperCAmelCase			:       str			     =Image.open(requests.get(__lowerCAmelCase						,    stream=__lowerCAmelCase  ).raw  ).convert('''RGB'''  )

      UpperCAmelCase			:       Tuple			     =SegformerImageProcessor()
      UpperCAmelCase			:       List[str]			     =processor(__lowerCAmelCase						,    return_tensors='''pt'''  ).pixel_values

      with torch.no_grad():
            UpperCAmelCase			:       Tuple			     =model(__lowerCAmelCase  )
            UpperCAmelCase			:       Optional[int]			     =outputs.logits

      print(logits.shape  )
      print('''First values of logits:'''						,    logits[0, 0, :3, :3]  )
      # assert values
      if model_name == "upernet-swin-tiny":
            UpperCAmelCase			:       str			     =torch.tensor(
                [[-7.5958, -7.5958, -7.4302], [-7.5958, -7.5958, -7.4302], [-7.4797, -7.4797, -7.3068]]  )
      elif model_name == "upernet-swin-small":
            UpperCAmelCase			:       List[str]			     =torch.tensor(
                [[-7.1921, -7.1921, -6.9532], [-7.1921, -7.1921, -6.9532], [-7.0908, -7.0908, -6.8534]]  )
      elif model_name == "upernet-swin-base":
            UpperCAmelCase			:       str			     =torch.tensor(
                [[-6.5851, -6.5851, -6.4330], [-6.5851, -6.5851, -6.4330], [-6.4763, -6.4763, -6.3254]]  )
      elif model_name == "upernet-swin-large":
            UpperCAmelCase			:       Tuple			     =torch.tensor(
                [[-7.5297, -7.5297, -7.3802], [-7.5297, -7.5297, -7.3802], [-7.4044, -7.4044, -7.2586]]  )
      print('''Logits:'''						,    outputs.logits[0, 0, :3, :3]  )
      assert torch.allclose(outputs.logits[0, 0, :3, :3]						,    __lowerCAmelCase						,    atol=1e-4  )
      print('''Looks ok!'''  )

      if pytorch_dump_folder_path is not None:
            print(f'''Saving model {model_name} to {pytorch_dump_folder_path}'''  )
            model.save_pretrained(__lowerCAmelCase  )
            print(f'''Saving processor to {pytorch_dump_folder_path}'''  )
            processor.save_pretrained(__lowerCAmelCase  )

      if push_to_hub:
            print(f'''Pushing model and processor for {model_name} to hub'''  )
            model.push_to_hub(f'''openmmlab/{model_name}'''  )
            processor.push_to_hub(f'''openmmlab/{model_name}'''  )


if __name__ == "__main__":
     __snake_case						=			argparse.ArgumentParser()
     # Required parameters
     parser.add_argument(
         '''--model_name''',
         default='''upernet-swin-tiny''',
         type=str,
         choices=[f'upernet-swin-{size}' for size in ['''tiny''', '''small''', '''base''', '''large''']],
         help='''Name of the Swin + UperNet 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 or not to push the converted model to the 🤗 hub.'''
     )

     __snake_case						=			parser.parse_args()
     convert_upernet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)








import gc
import unittest

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

from diffusers import HeunDiscreteScheduler, PriorTransformer, ShapEPipeline
from diffusers.pipelines.shap_e import ShapERenderer
from diffusers.utils import load_numpy, slow
from diffusers.utils.testing_utils import require_torch_gpu, torch_device

from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference






class 		lowercase_					(	lowercase ,  unittest.TestCase   ):






      '''simple docstring'''




      __snake_case       =  ShapEPipeline
      __snake_case       =  ['''prompt''']
      __snake_case       =  ['''prompt''']
      __snake_case       =  [
          '''num_images_per_prompt''',
          '''num_inference_steps''',
          '''generator''',
          '''latents''',
          '''guidance_scale''',
          '''frame_size''',
          '''output_type''',
          '''return_dict''',
      ]
      __snake_case       =  False







      @property
      def       __lowerCAmelCase					(      self      :						Union[str, Any]  )    ->Union[str, Any]:




       """simple docstring"""


       return 32







      @property
      def       __lowerCAmelCase					(      self      :						Optional[Any]  )    ->List[str]:




       """simple docstring"""


       return 32







      @property
      def       __lowerCAmelCase					(      self      :						Any  )    ->Tuple:




       """simple docstring"""


       return self.time_input_dim * 4







      @property
      def       __lowerCAmelCase					(      self      :						Tuple  )    ->Optional[Any]:




       """simple docstring"""


       return 8







      @property
      def       __lowerCAmelCase					(      self      :						Tuple  )    ->str:




       """simple docstring"""


       a		  =      CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''  )
       return tokenizer







      @property
      def       __lowerCAmelCase					(      self      :						Union[str, Any]  )    ->List[Any]:




       """simple docstring"""


       torch.manual_seed(0  )
       a		  =      CLIPTextConfig(
           bos_token_id=0 ,					eos_token_id=2 ,					hidden_size=self.text_embedder_hidden_size ,					projection_dim=self.text_embedder_hidden_size ,					intermediate_size=37 ,					layer_norm_eps=1e-0_5 ,					num_attention_heads=4 ,					num_hidden_layers=5 ,					pad_token_id=1 ,					vocab_size=1_000 ,					)
       return CLIPTextModelWithProjection(__UpperCAmelCase  )







      @property
      def       __lowerCAmelCase					(      self      :						Dict  )    ->Union[str, Any]:




       """simple docstring"""


       torch.manual_seed(0  )

       a		  =      {
           '''num_attention_heads''': 2,
           '''attention_head_dim''': 16,
           '''embedding_dim''': self.time_input_dim,
           '''num_embeddings''': 32,
           '''embedding_proj_dim''': self.text_embedder_hidden_size,
           '''time_embed_dim''': self.time_embed_dim,
           '''num_layers''': 1,
           '''clip_embed_dim''': self.time_input_dim * 2,
           '''additional_embeddings''': 0,
           '''time_embed_act_fn''': '''gelu''',
           '''norm_in_type''': '''layer''',
           '''encoder_hid_proj_type''': None,
           '''added_emb_type''': None,
       }

       a		  =      PriorTransformer(**__UpperCAmelCase  )
       return model







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




       """simple docstring"""


       torch.manual_seed(0  )

       a		  =      {
           '''param_shapes''': (
               (self.renderer_dim, 93),
               (self.renderer_dim, 8),
               (self.renderer_dim, 8),
               (self.renderer_dim, 8),
           ),
           '''d_latent''': self.time_input_dim,
           '''d_hidden''': self.renderer_dim,
           '''n_output''': 12,
           '''background''': (
               0.1,
               0.1,
               0.1,
           ),
       }
       a		  =      ShapERenderer(**__UpperCAmelCase  )
       return model







      def       __lowerCAmelCase					(      self      :						List[Any]  )    ->Any:




       """simple docstring"""


       a		  =      self.dummy_prior
       a		  =      self.dummy_text_encoder
       a		  =      self.dummy_tokenizer
       a		  =      self.dummy_renderer

       a		  =      HeunDiscreteScheduler(
           beta_schedule='''exp''' ,					num_train_timesteps=1_024 ,					prediction_type='''sample''' ,					use_karras_sigmas=__UpperCAmelCase ,					clip_sample=__UpperCAmelCase ,					clip_sample_range=1.0 ,					)
       a		  =      {
           '''prior''': prior,
           '''text_encoder''': text_encoder,
           '''tokenizer''': tokenizer,
           '''renderer''': renderer,
           '''scheduler''': scheduler,
       }

       return components







      def       __lowerCAmelCase					(      self      :						Tuple ,					__UpperCAmelCase      :						Optional[Any] ,					__UpperCAmelCase      :						str=0  )    ->Optional[int]:




       """simple docstring"""


       if str(__UpperCAmelCase  ).startswith('''mps'''  ):
        a		  =      torch.manual_seed(__UpperCAmelCase  )
       else:
        a		  =      torch.Generator(device=__UpperCAmelCase  ).manual_seed(__UpperCAmelCase  )
       a		  =      {
           '''prompt''': '''horse''',
           '''generator''': generator,
           '''num_inference_steps''': 1,
           '''frame_size''': 32,
           '''output_type''': '''np''',
       }
       return inputs







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




       """simple docstring"""


       a		  =      '''cpu'''

       a		  =      self.get_dummy_components()

       a		  =      self.pipeline_class(**__UpperCAmelCase  )
       a		  =      pipe.to(__UpperCAmelCase  )

       pipe.set_progress_bar_config(disable=__UpperCAmelCase  )

       a		  =      pipe(**self.get_dummy_inputs(__UpperCAmelCase  )  )
       a		  =      output.images[0]
       a		  =      image[0, -3:, -3:, -1]

       assert image.shape == (20, 32, 32, 3)

       a		  =      np.array(
           [
               0.00039216,
               0.00039216,
               0.00039216,
               0.00039216,
               0.00039216,
               0.00039216,
               0.00039216,
               0.00039216,
               0.00039216,
           ]  )

       assert np.abs(image_slice.flatten() - expected_slice  ).max() < 1e-2







      def       __lowerCAmelCase					(      self      :						Dict  )    ->Optional[Any]:




       """simple docstring"""


       self._test_inference_batch_consistent(batch_sizes=[1, 2]  )







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




       """simple docstring"""


       a		  =      torch_device == '''cpu'''
       a		  =      True

       self._test_inference_batch_single_identical(
           batch_size=2 ,					test_max_difference=__UpperCAmelCase ,					relax_max_difference=__UpperCAmelCase ,					)







      def       __lowerCAmelCase					(      self      :						str  )    ->Optional[int]:




       """simple docstring"""


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

       a		  =      1
       a		  =      2

       a		  =      self.get_dummy_inputs(__UpperCAmelCase  )

       for key in inputs.keys():
        if key in self.batch_params:
         a		  =      batch_size * [inputs[key]]

       a		  =      pipe(**__UpperCAmelCase ,					num_images_per_prompt=__UpperCAmelCase  )[0]

       assert images.shape[0] == batch_size * num_images_per_prompt






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






      '''simple docstring'''




      def       __lowerCAmelCase					(      self      :						int  )    ->Any:




       """simple docstring"""


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







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




       """simple docstring"""


       a		  =      load_numpy(
           '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
           '''/shap_e/test_shap_e_np_out.npy'''  )
       a		  =      ShapEPipeline.from_pretrained('''openai/shap-e'''  )
       a		  =      pipe.to(__UpperCAmelCase  )
       pipe.set_progress_bar_config(disable=__UpperCAmelCase  )

       a		  =      torch.Generator(device=__UpperCAmelCase  ).manual_seed(0  )

       a		  =      pipe(
           '''a shark''' ,					generator=__UpperCAmelCase ,					guidance_scale=15.0 ,					num_inference_steps=64 ,					frame_size=64 ,					output_type='''np''' ,					).images[0]

       assert images.shape == (20, 64, 64, 3)

       assert_mean_pixel_difference(__UpperCAmelCase ,					__UpperCAmelCase  )


