




import argparse
import re

import numpy as np
import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import (
    SamConfig,
    SamImageProcessor,
    SamModel,
    SamProcessor,
    SamVisionConfig,
)


SCREAMING_SNAKE_CASE__       :      Any  =   {
    'iou_prediction_head.layers.0': 'iou_prediction_head.proj_in',
    'iou_prediction_head.layers.1': 'iou_prediction_head.layers.0',
    'iou_prediction_head.layers.2': 'iou_prediction_head.proj_out',
    'mask_decoder.output_upscaling.0': 'mask_decoder.upscale_conv1',
    'mask_decoder.output_upscaling.1': 'mask_decoder.upscale_layer_norm',
    'mask_decoder.output_upscaling.3': 'mask_decoder.upscale_conv2',
    'mask_downscaling.0': 'mask_embed.conv1',
    'mask_downscaling.1': 'mask_embed.layer_norm1',
    'mask_downscaling.3': 'mask_embed.conv2',
    'mask_downscaling.4': 'mask_embed.layer_norm2',
    'mask_downscaling.6': 'mask_embed.conv3',
    'point_embeddings': 'point_embed',
    'pe_layer.positional_encoding_gaussian_matrix': 'shared_embedding.positional_embedding',
    'image_encoder': 'vision_encoder',
    'neck.0': 'neck.conv1',
    'neck.1': 'neck.layer_norm1',
    'neck.2': 'neck.conv2',
    'neck.3': 'neck.layer_norm2',
    'patch_embed.proj': 'patch_embed.projection',
    '.norm': '.layer_norm',
    'blocks': 'layers',
}

def      A      (      _SCREAMING_SNAKE_CASE    )	-> Tuple:
       lowerCamelCase   :      Optional[int]			 =					{}
       state_dict.pop("pixel_mean" ,_SCREAMING_SNAKE_CASE    )
       state_dict.pop("pixel_std" ,_SCREAMING_SNAKE_CASE    )

       lowerCamelCase   :      Tuple			 =					r".*.output_hypernetworks_mlps.(\d+).layers.(\d+).*"

       for key, value in state_dict.items():
              for key_to_modify, new_key in KEYS_TO_MODIFY_MAPPING.items():
                     if key_to_modify in key:
                            lowerCamelCase   :      Dict			 =					key.replace(_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE    )

              if re.match(_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE    ):
                     lowerCamelCase   :      str			 =					int(re.match(_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE    ).group(2    )    )
                     if layer_nb == 0:
                            lowerCamelCase   :      int			 =					key.replace("layers.0" ,"proj_in"    )
                     elif layer_nb == 1:
                            lowerCamelCase   :      int			 =					key.replace("layers.1" ,"layers.0"    )
                     elif layer_nb == 2:
                            lowerCamelCase   :      Tuple			 =					key.replace("layers.2" ,"proj_out"    )

              lowerCamelCase   :      List[Any]			 =					value

       lowerCamelCase   :      Tuple			 =					model_state_dict[
           "prompt_encoder.shared_embedding.positional_embedding"
       ]

       return model_state_dict

def      A      (      _SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE="ybelkada/segment-anything"    )	-> List[str]:
       lowerCamelCase   :      Union[str, Any]			 =					hf_hub_download(_SCREAMING_SNAKE_CASE ,f'''checkpoints/{model_name}.pth'''    )

       if "sam_vit_b" in model_name:
              lowerCamelCase   :      Optional[int]			 =					SamConfig()
       elif "sam_vit_l" in model_name:
              lowerCamelCase   :      Optional[Any]			 =					SamVisionConfig(
                  hidden_size=1024 ,num_hidden_layers=24 ,num_attention_heads=16 ,global_attn_indexes=[5, 11, 17, 23] ,)

              lowerCamelCase   :      int			 =					SamConfig(
                  vision_config=_SCREAMING_SNAKE_CASE ,)
       elif "sam_vit_h" in model_name:
              lowerCamelCase   :      Optional[int]			 =					SamVisionConfig(
                  hidden_size=1280 ,num_hidden_layers=32 ,num_attention_heads=16 ,global_attn_indexes=[7, 15, 23, 31] ,)

              lowerCamelCase   :      int			 =					SamConfig(
                  vision_config=_SCREAMING_SNAKE_CASE ,)

       lowerCamelCase   :      Optional[Any]			 =					torch.load(_SCREAMING_SNAKE_CASE ,map_location="cpu"    )
       lowerCamelCase   :      List[Any]			 =					replace_keys(_SCREAMING_SNAKE_CASE    )

       lowerCamelCase   :      Optional[int]			 =					SamImageProcessor()

       lowerCamelCase   :      List[str]			 =					SamProcessor(image_processor=_SCREAMING_SNAKE_CASE    )
       lowerCamelCase   :      List[Any]			 =					SamModel(_SCREAMING_SNAKE_CASE    )

       hf_model.load_state_dict(_SCREAMING_SNAKE_CASE    )
       lowerCamelCase   :      Optional[int]			 =					hf_model.to("cuda"    )

       lowerCamelCase   :      Dict			 =					"https://huggingface.co/ybelkada/segment-anything/resolve/main/assets/car.png"
       lowerCamelCase   :      Tuple			 =					Image.open(requests.get(_SCREAMING_SNAKE_CASE ,stream=_SCREAMING_SNAKE_CASE    ).raw    ).convert("RGB"    )

       lowerCamelCase   :      Dict			 =					[[[400, 650]]]
       lowerCamelCase   :      str			 =					[[1]]

       lowerCamelCase   :      Dict			 =					processor(images=np.array(_SCREAMING_SNAKE_CASE    ) ,return_tensors="pt"    ).to("cuda"    )

       with torch.no_grad():
              lowerCamelCase   :      str			 =					hf_model(**_SCREAMING_SNAKE_CASE    )
       lowerCamelCase   :      Any			 =					output.iou_scores.squeeze()

       if model_name == "sam_vit_h_4b8939":
              assert scores[-1].item() == 0.579890251159668

              lowerCamelCase   :      str			 =					processor(
                  images=np.array(_SCREAMING_SNAKE_CASE    ) ,input_points=_SCREAMING_SNAKE_CASE ,input_labels=_SCREAMING_SNAKE_CASE ,return_tensors="pt"    ).to("cuda"    )

              with torch.no_grad():
                     lowerCamelCase   :      Any			 =					hf_model(**_SCREAMING_SNAKE_CASE    )
              lowerCamelCase   :      List[str]			 =					output.iou_scores.squeeze()

              assert scores[-1].item() == 0.9712603092193604

              lowerCamelCase   :      Any			 =					((75, 275, 1725, 850),)

              lowerCamelCase   :      Any			 =					processor(images=np.array(_SCREAMING_SNAKE_CASE    ) ,input_boxes=_SCREAMING_SNAKE_CASE ,return_tensors="pt"    ).to("cuda"    )

              with torch.no_grad():
                     lowerCamelCase   :      Union[str, Any]			 =					hf_model(**_SCREAMING_SNAKE_CASE    )
              lowerCamelCase   :      int			 =					output.iou_scores.squeeze()

              assert scores[-1].item() == 0.8686015605926514

              # Test with 2 points and 1 image.
              lowerCamelCase   :      Optional[Any]			 =					[[[400, 650], [800, 650]]]
              lowerCamelCase   :      Any			 =					[[1, 1]]

              lowerCamelCase   :      int			 =					processor(
                  images=np.array(_SCREAMING_SNAKE_CASE    ) ,input_points=_SCREAMING_SNAKE_CASE ,input_labels=_SCREAMING_SNAKE_CASE ,return_tensors="pt"    ).to("cuda"    )

              with torch.no_grad():
                     lowerCamelCase   :      Optional[Any]			 =					hf_model(**_SCREAMING_SNAKE_CASE    )
              lowerCamelCase   :      int			 =					output.iou_scores.squeeze()

              assert scores[-1].item() == 0.9936047792434692


if __name__ == "__main__":
       SCREAMING_SNAKE_CASE__       :      str  =   argparse.ArgumentParser()
       SCREAMING_SNAKE_CASE__       :      int  =   ['sam_vit_b_01ec64', 'sam_vit_h_4b8939', 'sam_vit_l_0b3195']
       parser.add_argument(
           '--model_name',
           default='sam_vit_h_4b8939',
           choices=choices,
           type=str,
           help='Path to hf config.json of model to convert',
       )
       parser.add_argument('--pytorch_dump_folder_path', default=None, type=str, help='Path to the output PyTorch model.')
       parser.add_argument(
           '--push_to_hub',
           action='store_true',
           help='Whether to push the model and processor to the hub after converting',
       )
       parser.add_argument(
           '--model_hub_id',
           default='ybelkada/segment-anything',
           choices=choices,
           type=str,
           help='Path to hf config.json of model to convert',
       )

       SCREAMING_SNAKE_CASE__       :      Any  =   parser.parse_args()

       convert_sam_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub, args.model_hub_id)











"""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 							_UpperCAmelCase							( a		,a		,a		,unittest.TestCase      ):





     '''simple docstring'''



     a__									=StableDiffusionAttendAndExcitePipeline
     a__									=False
     a__									=TEXT_TO_IMAGE_PARAMS
     a__									=TEXT_TO_IMAGE_BATCH_PARAMS.union({'''token_indices'''}      )
     a__									=TEXT_TO_IMAGE_IMAGE_PARAMS
     a__									=TEXT_TO_IMAGE_IMAGE_PARAMS





     @classmethod
     def  __lowerCAmelCase	(					cls		)							->					List[str]:
         super().setUpClass()
         torch.use_deterministic_algorithms(A		)





     @classmethod
     def  __lowerCAmelCase	(					cls		)							->					Union[str, Any]:
         super().tearDownClass()
         torch.use_deterministic_algorithms(A		)





     def  __lowerCAmelCase	(					self		)							->					Tuple:
         torch.manual_seed(0		)
         _UpperCAmelCase							:  Optional[int]											=    UNetaDConditionModel(
             block_out_channels=(3_2, 6_4)       ,	layers_per_block=1       ,	sample_size=3_2       ,	in_channels=4       ,	out_channels=4       ,	down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')       ,	up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''')       ,	cross_attention_dim=3_2       ,	attention_head_dim=(2, 4)       ,	use_linear_projection=A       ,	)
         _UpperCAmelCase							:  List[Any]											=    DDIMScheduler(
             beta_start=0.00_085       ,	beta_end=0.012       ,	beta_schedule='''scaled_linear'''       ,	clip_sample=A       ,	set_alpha_to_one=A       ,	)
         torch.manual_seed(0		)
         _UpperCAmelCase							:  int											=    AutoencoderKL(
             block_out_channels=[3_2, 6_4]       ,	in_channels=3       ,	out_channels=3       ,	down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']       ,	up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']       ,	latent_channels=4       ,	sample_size=1_2_8       ,	)
         torch.manual_seed(0		)
         _UpperCAmelCase							:  int											=    CLIPTextConfig(
             bos_token_id=0       ,	eos_token_id=2       ,	hidden_size=3_2       ,	intermediate_size=3_7       ,	layer_norm_eps=1E-05       ,	num_attention_heads=4       ,	num_hidden_layers=5       ,	pad_token_id=1       ,	vocab_size=1_0_0_0       ,	hidden_act='''gelu'''       ,	projection_dim=5_1_2       ,	)
         _UpperCAmelCase							:  List[str]											=    CLIPTextModel(A		)
         _UpperCAmelCase							:  str											=    CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''		)

         _UpperCAmelCase							:  Union[str, Any]											=    {
             '''unet''': unet,
             '''scheduler''': scheduler,
             '''vae''': vae,
             '''text_encoder''': text_encoder,
             '''tokenizer''': tokenizer,
             '''safety_checker''': None,
             '''feature_extractor''': None,
         }

         return components





     def  __lowerCAmelCase	(					self       ,	A       ,	A=0		)							->					List[Any]:
         if str(A		).startswith('''mps'''		):
             _UpperCAmelCase							:  Optional[int]											=    torch.manual_seed(A		)
         else:
             _UpperCAmelCase							:  Union[str, Any]											=    torch.Generator(device=A		).manual_seed(A		)
         _UpperCAmelCase							:  List[str]											=    {
             '''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  __lowerCAmelCase	(					self		)							->					int:
         _UpperCAmelCase							:  List[str]											=    '''cpu'''

         _UpperCAmelCase							:  Tuple											=    self.get_dummy_components()
         _UpperCAmelCase							:  int											=    self.pipeline_class(**A		)
         pipe.to(A		)
         pipe.set_progress_bar_config(disable=A		)

         _UpperCAmelCase							:  Dict											=    self.get_dummy_inputs(A		)
         _UpperCAmelCase							:  Union[str, Any]											=    pipe(**A		).images
         _UpperCAmelCase							:  Tuple											=    image[0, -3:, -3:, -1]

         self.assertEqual(image.shape       ,	(1, 6_4, 6_4, 3)		)
         _UpperCAmelCase							:  int											=    np.array(
             [0.63_905_364, 0.62_897_307, 0.48_599_017, 0.5_133_624, 0.5_550_048, 0.45_769_516, 0.50_326_973, 0.5_023_139, 0.45_384_496]		)
         _UpperCAmelCase							:  Tuple											=    np.abs(image_slice.flatten() - expected_slice		).max()
         self.assertLessEqual(A       ,	1E-3		)





     def  __lowerCAmelCase	(					self		)							->					Dict:
         super().test_cpu_offload_forward_pass(expected_max_diff=5E-4		)





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





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





     def  __lowerCAmelCase	(					self		)							->					List[str]:
         super().test_dict_tuple_outputs_equivalent(expected_max_difference=3E-3		)





     def  __lowerCAmelCase	(					self		)							->					List[str]:
         super().test_pt_np_pil_outputs_equivalent(expected_max_diff=5E-4		)





     def  __lowerCAmelCase	(					self		)							->					str:
         super().test_save_load_local(expected_max_difference=5E-4		)





     def  __lowerCAmelCase	(					self		)							->					Optional[int]:
         super().test_save_load_optional_components(expected_max_difference=4E-4		)



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





     '''simple docstring'''



     @classmethod
     def  __lowerCAmelCase	(					cls		)							->					Union[str, Any]:
         super().setUpClass()
         torch.use_deterministic_algorithms(A		)





     @classmethod
     def  __lowerCAmelCase	(					cls		)							->					Optional[int]:
         super().tearDownClass()
         torch.use_deterministic_algorithms(A		)





     def  __lowerCAmelCase	(					self		)							->					List[str]:
         super().tearDown()
         gc.collect()
         torch.cuda.empty_cache()





     def  __lowerCAmelCase	(					self		)							->					str:
         _UpperCAmelCase							:  Any											=    torch.manual_seed(5_1		)

         _UpperCAmelCase							:  Optional[Any]											=    StableDiffusionAttendAndExcitePipeline.from_pretrained(
             '''CompVis/stable-diffusion-v1-4'''       ,	safety_checker=A       ,	torch_dtype=torch.floataa		)
         pipe.to('''cuda'''		)

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

         _UpperCAmelCase							:  Dict											=    pipe(
             prompt=A       ,	token_indices=A       ,	guidance_scale=7.5       ,	generator=A       ,	num_inference_steps=5       ,	max_iter_to_alter=5       ,	output_type='''numpy'''       ,	).images[0]

         _UpperCAmelCase							:  List[Any]											=    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


