

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


_lowerCamelCase    :       List[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__					(		UpperCAmelCase							:      Optional[int]      )  ->						List[Any]:
  UpperCAmelCase	:							List[str]						=			{}
  state_dict.pop('''pixel_mean'''    ,			UpperCAmelCase      )
  state_dict.pop('''pixel_std'''    ,			UpperCAmelCase      )

  UpperCAmelCase	:							Dict						=			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:
        UpperCAmelCase	:							List[str]						=			key.replace(UpperCAmelCase    ,			UpperCAmelCase      )

    if re.match(UpperCAmelCase    ,			UpperCAmelCase      ):
      UpperCAmelCase	:							List[Any]						=			int(re.match(UpperCAmelCase    ,			UpperCAmelCase      ).group(2      )      )
      if layer_nb == 0:
        UpperCAmelCase	:							Union[str, Any]						=			key.replace('''layers.0'''    ,			'''proj_in'''      )
      elif layer_nb == 1:
        UpperCAmelCase	:							Optional[int]						=			key.replace('''layers.1'''    ,			'''layers.0'''      )
      elif layer_nb == 2:
        UpperCAmelCase	:							Dict						=			key.replace('''layers.2'''    ,			'''proj_out'''      )

    UpperCAmelCase	:							int						=			value

  UpperCAmelCase	:							Any						=			model_state_dict[
      '''prompt_encoder.shared_embedding.positional_embedding'''
  ]

  return model_state_dict

def       a__					(		UpperCAmelCase							:      Dict    ,			UpperCAmelCase							:      Dict    ,			UpperCAmelCase							:      int    ,			UpperCAmelCase							:      Tuple="ybelkada/segment-anything"      )  ->						Optional[int]:
  UpperCAmelCase	:							Optional[Any]						=			hf_hub_download(UpperCAmelCase    ,			f'''checkpoints/{model_name}.pth'''      )

  if "sam_vit_b" in model_name:
    UpperCAmelCase	:							List[Any]						=			SamConfig()
  elif "sam_vit_l" in model_name:
    UpperCAmelCase	:							List[Any]						=			SamVisionConfig(
        hidden_size=1_024    ,			num_hidden_layers=24    ,			num_attention_heads=16    ,			global_attn_indexes=[5, 11, 17, 23]    ,			)

    UpperCAmelCase	:							Any						=			SamConfig(
        vision_config=UpperCAmelCase    ,			)
  elif "sam_vit_h" in model_name:
    UpperCAmelCase	:							Optional[int]						=			SamVisionConfig(
        hidden_size=1_280    ,			num_hidden_layers=32    ,			num_attention_heads=16    ,			global_attn_indexes=[7, 15, 23, 31]    ,			)

    UpperCAmelCase	:							List[Any]						=			SamConfig(
        vision_config=UpperCAmelCase    ,			)

  UpperCAmelCase	:							Dict						=			torch.load(UpperCAmelCase    ,			map_location='''cpu'''      )
  UpperCAmelCase	:							Union[str, Any]						=			replace_keys(UpperCAmelCase      )

  UpperCAmelCase	:							List[Any]						=			SamImageProcessor()

  UpperCAmelCase	:							Union[str, Any]						=			SamProcessor(image_processor=UpperCAmelCase      )
  UpperCAmelCase	:							Union[str, Any]						=			SamModel(UpperCAmelCase      )

  hf_model.load_state_dict(UpperCAmelCase      )
  UpperCAmelCase	:							Dict						=			hf_model.to('''cuda'''      )

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

  UpperCAmelCase	:							Dict						=			[[[400, 650]]]
  UpperCAmelCase	:							List[str]						=			[[1]]

  UpperCAmelCase	:							Optional[int]						=			processor(images=np.array(UpperCAmelCase      )    ,			return_tensors='''pt'''      ).to('''cuda'''      )

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

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

    UpperCAmelCase	:							str						=			processor(
        images=np.array(UpperCAmelCase      )    ,			input_points=UpperCAmelCase    ,			input_labels=UpperCAmelCase    ,			return_tensors='''pt'''      ).to('''cuda'''      )

    with torch.no_grad():
      UpperCAmelCase	:							Optional[Any]						=			hf_model(**UpperCAmelCase      )
    UpperCAmelCase	:							Union[str, Any]						=			output.iou_scores.squeeze()

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

    UpperCAmelCase	:							Dict						=			((75, 275, 1_725, 850),)

    UpperCAmelCase	:							str						=			processor(images=np.array(UpperCAmelCase      )    ,			input_boxes=UpperCAmelCase    ,			return_tensors='''pt'''      ).to('''cuda'''      )

    with torch.no_grad():
      UpperCAmelCase	:							Any						=			hf_model(**UpperCAmelCase      )
    UpperCAmelCase	:							Dict						=			output.iou_scores.squeeze()

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

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

    UpperCAmelCase	:							Union[str, Any]						=			processor(
        images=np.array(UpperCAmelCase      )    ,			input_points=UpperCAmelCase    ,			input_labels=UpperCAmelCase    ,			return_tensors='''pt'''      ).to('''cuda'''      )

    with torch.no_grad():
      UpperCAmelCase	:							int						=			hf_model(**UpperCAmelCase      )
    UpperCAmelCase	:							List[Any]						=			output.iou_scores.squeeze()

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


if __name__ == "__main__":
     _lowerCamelCase    :       Union[str, Any]        =							argparse.ArgumentParser()
     _lowerCamelCase    :       List[Any]        =							["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",
     )

     _lowerCamelCase    :       Optional[Any]        =							parser.parse_args()

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








import gc
import random
import unittest

import numpy as np
import torch
from transformers import CLIPImageProcessor, CLIPVisionConfig, CLIPVisionModel

from diffusers import HeunDiscreteScheduler, PriorTransformer, ShapEImgaImgPipeline
from diffusers.pipelines.shap_e import ShapERenderer
from diffusers.utils import floats_tensor, load_image, 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  __UpperCAmelCase					(	lowerCamelCase__      ,  unittest.TestCase						):
   UpperCamelCase       					=					ShapEImgaImgPipeline
   UpperCamelCase       					=					["""image"""]
   UpperCamelCase       					=					["""image"""]
   UpperCamelCase       					=					[
       """num_images_per_prompt""",
       """num_inference_steps""",
       """generator""",
       """latents""",
       """guidance_scale""",
       """frame_size""",
       """output_type""",
       """return_dict""",
   ]
   UpperCamelCase       					=					False



   @property
   def 	__magic_name__   (							self							:					Optional[Any]						):
     return 3_2



   @property
   def 	__magic_name__   (							self							:					Optional[int]						):
     return 3_2



   @property
   def 	__magic_name__   (							self							:					Union[str, Any]						):
     return self.time_input_dim * 4



   @property
   def 	__magic_name__   (							self							:					Any						):
     return 8



   @property
   def 	__magic_name__   (							self							:					List[str]						):
     torch.manual_seed(0						)
     UpperCAmelCase	:							str						=			CLIPVisionConfig(
         hidden_size=self.text_embedder_hidden_size,		image_size=6_4,		projection_dim=self.text_embedder_hidden_size,		intermediate_size=3_7,		num_attention_heads=4,		num_channels=3,		num_hidden_layers=5,		patch_size=1,		)

     UpperCAmelCase	:							Any						=			CLIPVisionModel(__A						)
     return model



   @property
   def 	__magic_name__   (							self							:					int						):
     UpperCAmelCase	:							Optional[int]						=			CLIPImageProcessor(
         crop_size=2_2_4,		do_center_crop=__A,		do_normalize=__A,		do_resize=__A,		image_mean=[0.4_8_1_4_5_4_6_6, 0.4_5_7_8_2_7_5, 0.4_0_8_2_1_0_7_3],		image_std=[0.2_6_8_6_2_9_5_4, 0.2_6_1_3_0_2_5_8, 0.2_7_5_7_7_7_1_1],		resample=3,		size=2_2_4,		)

     return image_processor



   @property
   def 	__magic_name__   (							self							:					Dict						):
     torch.manual_seed(0						)

     UpperCAmelCase	:							Any						=			{
         '''num_attention_heads''': 2,
         '''attention_head_dim''': 1_6,
         '''embedding_dim''': self.time_input_dim,
         '''num_embeddings''': 3_2,
         '''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''',
         '''embedding_proj_norm_type''': '''layer''',
         '''encoder_hid_proj_type''': None,
         '''added_emb_type''': None,
     }

     UpperCAmelCase	:							List[Any]						=			PriorTransformer(**__A						)
     return model



   @property
   def 	__magic_name__   (							self							:					Union[str, Any]						):
     torch.manual_seed(0						)

     UpperCAmelCase	:							List[Any]						=			{
         '''param_shapes''': (
             (self.renderer_dim, 9_3),
             (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''': 1_2,
         '''background''': (
             0.1,
             0.1,
             0.1,
         ),
     }
     UpperCAmelCase	:							List[str]						=			ShapERenderer(**__A						)
     return model



   def 	__magic_name__   (							self							:					List[Any]						):
     UpperCAmelCase	:							str						=			self.dummy_prior
     UpperCAmelCase	:							List[str]						=			self.dummy_image_encoder
     UpperCAmelCase	:							List[Any]						=			self.dummy_image_processor
     UpperCAmelCase	:							Dict						=			self.dummy_renderer

     UpperCAmelCase	:							int						=			HeunDiscreteScheduler(
         beta_schedule='''exp''',		num_train_timesteps=1_0_2_4,		prediction_type='''sample''',		use_karras_sigmas=__A,		clip_sample=__A,		clip_sample_range=1.0,		)
     UpperCAmelCase	:							List[str]						=			{
         '''prior''': prior,
         '''image_encoder''': image_encoder,
         '''image_processor''': image_processor,
         '''renderer''': renderer,
         '''scheduler''': scheduler,
     }

     return components



   def 	__magic_name__   (							self							:					List[Any],		__A							:					Dict,		__A							:					List[Any]=0						):
     UpperCAmelCase	:							int						=			floats_tensor((1, 3, 6_4, 6_4),		rng=random.Random(__A						)						).to(__A						)

     if str(__A						).startswith('''mps'''						):
       UpperCAmelCase	:							List[str]						=			torch.manual_seed(__A						)
     else:
       UpperCAmelCase	:							Optional[Any]						=			torch.Generator(device=__A						).manual_seed(__A						)
     UpperCAmelCase	:							Dict						=			{
         '''image''': input_image,
         '''generator''': generator,
         '''num_inference_steps''': 1,
         '''frame_size''': 3_2,
         '''output_type''': '''np''',
     }
     return inputs



   def 	__magic_name__   (							self							:					str						):
     UpperCAmelCase	:							Dict						=			'''cpu'''

     UpperCAmelCase	:							Any						=			self.get_dummy_components()

     UpperCAmelCase	:							Tuple						=			self.pipeline_class(**__A						)
     UpperCAmelCase	:							Dict						=			pipe.to(__A						)

     pipe.set_progress_bar_config(disable=__A						)

     UpperCAmelCase	:							List[Any]						=			pipe(**self.get_dummy_inputs(__A						)						)
     UpperCAmelCase	:							Any						=			output.images[0]
     UpperCAmelCase	:							List[str]						=			image[0, -3:, -3:, -1]

     assert image.shape == (2_0, 3_2, 3_2, 3)

     UpperCAmelCase	:							Optional[int]						=			np.array(
         [
             0.0_0_0_3_9_2_1_6,
             0.0_0_0_3_9_2_1_6,
             0.0_0_0_3_9_2_1_6,
             0.0_0_0_3_9_2_1_6,
             0.0_0_0_3_9_2_1_6,
             0.0_0_0_3_9_2_1_6,
             0.0_0_0_3_9_2_1_6,
             0.0_0_0_3_9_2_1_6,
             0.0_0_0_3_9_2_1_6,
         ]						)

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



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



   def 	__magic_name__   (							self							:					str						):
     UpperCAmelCase	:							Tuple						=			torch_device == '''cpu'''
     UpperCAmelCase	:							List[Any]						=			True
     self._test_inference_batch_single_identical(
         batch_size=2,		test_max_difference=__A,		relax_max_difference=__A,		)





   def 	__magic_name__   (							self							:					Dict						):
     UpperCAmelCase	:							Dict						=			self.get_dummy_components()
     UpperCAmelCase	:							Union[str, Any]						=			self.pipeline_class(**__A						)
     UpperCAmelCase	:							List[str]						=			pipe.to(__A						)
     pipe.set_progress_bar_config(disable=__A						)

     UpperCAmelCase	:							Any						=			1
     UpperCAmelCase	:							Any						=			2

     UpperCAmelCase	:							Optional[int]						=			self.get_dummy_inputs(__A						)

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

     UpperCAmelCase	:							int						=			pipe(**__A,		num_images_per_prompt=__A						)[0]

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






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



   def 	__magic_name__   (							self							:					List[Any]						):
     # clean up the VRAM after each test
     super().tearDown()
     gc.collect()
     torch.cuda.empty_cache()





   def 	__magic_name__   (							self							:					Tuple						):
     UpperCAmelCase	:							List[str]						=			load_image(
         '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main''' '''/shap_e/corgi.png'''						)
     UpperCAmelCase	:							Dict						=			load_numpy(
         '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
         '''/shap_e/test_shap_e_img2img_out.npy'''						)
     UpperCAmelCase	:							Union[str, Any]						=			ShapEImgaImgPipeline.from_pretrained('''openai/shap-e-img2img'''						)
     UpperCAmelCase	:							int						=			pipe.to(__A						)
     pipe.set_progress_bar_config(disable=__A						)

     UpperCAmelCase	:							Dict						=			torch.Generator(device=__A						).manual_seed(0						)

     UpperCAmelCase	:							List[str]						=			pipe(
         __A,		generator=__A,		guidance_scale=3.0,		num_inference_steps=6_4,		frame_size=6_4,		output_type='''np''',		).images[0]

     assert images.shape == (2_0, 6_4, 6_4, 3)

     assert_mean_pixel_difference(__A,		__A						)





