



from typing import Any, Dict, Optional

import torch
import torch.nn.functional as F
from torch import nn

from ..utils import maybe_allow_in_graph
from .activations import get_activation
from .attention_processor import Attention
from .embeddings import CombinedTimestepLabelEmbeddings



@maybe_allow_in_graph
class 			__lowerCAmelCase  (						nn.Module ):




 def __init__(    self						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase=0.0						,					__UpperCAmelCase = None						,					__UpperCAmelCase = "geglu"						,					__UpperCAmelCase = None						,					__UpperCAmelCase = False						,					__UpperCAmelCase = False						,					__UpperCAmelCase = False						,					__UpperCAmelCase = False						,					__UpperCAmelCase = True						,					__UpperCAmelCase = "layer_norm"						,					__UpperCAmelCase = False						,					):



    '''simple docstring'''
    super().__init__()
    __lowerCamelCase    				=     only_cross_attention

    __lowerCamelCase    				=     (num_embeds_ada_norm is not None) and norm_type == '''ada_norm_zero'''
    __lowerCamelCase    				=     (num_embeds_ada_norm is not None) and norm_type == '''ada_norm'''

    if norm_type in ("ada_norm", "ada_norm_zero") and num_embeds_ada_norm is None:
       raise ValueError(
           F"""`norm_type` is set to {norm_type}, but `num_embeds_ada_norm` is not defined. Please make sure to"""
           F""" define `num_embeds_ada_norm` if setting `norm_type` to {norm_type}."""					)

    # Define 3 blocks. Each block has its own normalization layer.
    # 1. Self-Attn
    if self.use_ada_layer_norm:
       __lowerCamelCase    				=     AdaLayerNorm(__UpperCAmelCase						,					__UpperCAmelCase					)
    elif self.use_ada_layer_norm_zero:
       __lowerCamelCase    				=     AdaLayerNormZero(__UpperCAmelCase						,					__UpperCAmelCase					)
    else:
       __lowerCamelCase    				=     nn.LayerNorm(__UpperCAmelCase						,					elementwise_affine=__UpperCAmelCase					)
    __lowerCamelCase    				=     Attention(
        query_dim=__UpperCAmelCase						,					heads=__UpperCAmelCase						,					dim_head=__UpperCAmelCase						,					dropout=__UpperCAmelCase						,					bias=__UpperCAmelCase						,					cross_attention_dim=cross_attention_dim if only_cross_attention else None						,					upcast_attention=__UpperCAmelCase						,					)

    # 2. Cross-Attn
    if cross_attention_dim is not None or double_self_attention:
       # We currently only use AdaLayerNormZero for self attention where there will only be one attention block.
       # I.e. the number of returned modulation chunks from AdaLayerZero would not make sense if returned during
       # the second cross attention block.
       __lowerCamelCase    				=     (
           AdaLayerNorm(__UpperCAmelCase						,					__UpperCAmelCase					)
           if self.use_ada_layer_norm
           else nn.LayerNorm(__UpperCAmelCase						,					elementwise_affine=__UpperCAmelCase					)
       )
       __lowerCamelCase    				=     Attention(
           query_dim=__UpperCAmelCase						,					cross_attention_dim=cross_attention_dim if not double_self_attention else None						,					heads=__UpperCAmelCase						,					dim_head=__UpperCAmelCase						,					dropout=__UpperCAmelCase						,					bias=__UpperCAmelCase						,					upcast_attention=__UpperCAmelCase						,					)  # is self-attn if encoder_hidden_states is none
    else:
       __lowerCamelCase    				=     None
       __lowerCamelCase    				=     None

    # 3. Feed-forward
    __lowerCamelCase    				=     nn.LayerNorm(__UpperCAmelCase						,					elementwise_affine=__UpperCAmelCase					)
    __lowerCamelCase    				=     FeedForward(__UpperCAmelCase						,					dropout=__UpperCAmelCase						,					activation_fn=__UpperCAmelCase						,					final_dropout=__UpperCAmelCase					)

    # let chunk size default to None
    __lowerCamelCase    				=     None
    __lowerCamelCase    				=     0




 def 							lowerCamelCase   (    self						,					__UpperCAmelCase						,					__UpperCAmelCase					):



    '''simple docstring'''
    # Sets chunk feed-forward
    __lowerCamelCase    				=     chunk_size
    __lowerCamelCase    				=     dim




 def 							lowerCamelCase   (    self						,					__UpperCAmelCase						,					__UpperCAmelCase = None						,					__UpperCAmelCase = None						,					__UpperCAmelCase = None						,					__UpperCAmelCase = None						,					__UpperCAmelCase = None						,					__UpperCAmelCase = None						,					):



    '''simple docstring'''
    # Notice that normalization is always applied before the real computation in the following blocks.
    # 1. Self-Attention
    if self.use_ada_layer_norm:
       __lowerCamelCase    				=     self.norma(__UpperCAmelCase						,					__UpperCAmelCase					)
    elif self.use_ada_layer_norm_zero:
       __lowerCamelCase				,__lowerCamelCase				,__lowerCamelCase				,__lowerCamelCase				,__lowerCamelCase    				=     self.norma(
           __UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase						,					hidden_dtype=hidden_states.dtype					)
    else:
       __lowerCamelCase    				=     self.norma(__UpperCAmelCase					)

    __lowerCamelCase    				=     cross_attention_kwargs if cross_attention_kwargs is not None else {}

    __lowerCamelCase    				=     self.attna(
        __UpperCAmelCase						,					encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None						,					attention_mask=__UpperCAmelCase						,					**__UpperCAmelCase						,					)
    if self.use_ada_layer_norm_zero:
       __lowerCamelCase    				=     gate_msa.unsqueeze(1					) * attn_output
    __lowerCamelCase    				=     attn_output + hidden_states

    # 2. Cross-Attention
    if self.attna is not None:
       __lowerCamelCase    				=     (
           self.norma(__UpperCAmelCase						,					__UpperCAmelCase					) if self.use_ada_layer_norm else self.norma(__UpperCAmelCase					)
       )

       __lowerCamelCase    				=     self.attna(
           __UpperCAmelCase						,					encoder_hidden_states=__UpperCAmelCase						,					attention_mask=__UpperCAmelCase						,					**__UpperCAmelCase						,					)
       __lowerCamelCase    				=     attn_output + hidden_states

    # 3. Feed-forward
    __lowerCamelCase    				=     self.norma(__UpperCAmelCase					)

    if self.use_ada_layer_norm_zero:
       __lowerCamelCase    				=     norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]

    if self._chunk_size is not None:
       # "feed_forward_chunk_size" can be used to save memory
       if norm_hidden_states.shape[self._chunk_dim] % self._chunk_size != 0:
          raise ValueError(
              F"""`hidden_states` dimension to be chunked: {norm_hidden_states.shape[self._chunk_dim]} has to be divisible by chunk size: {self._chunk_size}. Make sure to set an appropriate `chunk_size` when calling `unet.enable_forward_chunking`."""					)

       __lowerCamelCase    				=     norm_hidden_states.shape[self._chunk_dim] // self._chunk_size
       __lowerCamelCase    				=     torch.cat(
           [self.ff(__UpperCAmelCase					) for hid_slice in norm_hidden_states.chunk(__UpperCAmelCase						,					dim=self._chunk_dim					)]						,					dim=self._chunk_dim						,					)
    else:
       __lowerCamelCase    				=     self.ff(__UpperCAmelCase					)

    if self.use_ada_layer_norm_zero:
       __lowerCamelCase    				=     gate_mlp.unsqueeze(1					) * ff_output

    __lowerCamelCase    				=     ff_output + hidden_states

    return hidden_states



class 			__lowerCAmelCase  (						nn.Module ):




 def __init__(    self						,					__UpperCAmelCase						,					__UpperCAmelCase = None						,					__UpperCAmelCase = 4						,					__UpperCAmelCase = 0.0						,					__UpperCAmelCase = "geglu"						,					__UpperCAmelCase = False						,					):



    '''simple docstring'''
    super().__init__()
    __lowerCamelCase    				=     int(dim * mult					)
    __lowerCamelCase    				=     dim_out if dim_out is not None else dim

    if activation_fn == "gelu":
       __lowerCamelCase    				=     GELU(__UpperCAmelCase						,					__UpperCAmelCase					)
    if activation_fn == "gelu-approximate":
       __lowerCamelCase    				=     GELU(__UpperCAmelCase						,					__UpperCAmelCase						,					approximate='''tanh'''					)
    elif activation_fn == "geglu":
       __lowerCamelCase    				=     GEGLU(__UpperCAmelCase						,					__UpperCAmelCase					)
    elif activation_fn == "geglu-approximate":
       __lowerCamelCase    				=     ApproximateGELU(__UpperCAmelCase						,					__UpperCAmelCase					)

    __lowerCamelCase    				=     nn.ModuleList([]					)
    # project in
    self.net.append(__UpperCAmelCase					)
    # project dropout
    self.net.append(nn.Dropout(__UpperCAmelCase					)					)
    # project out
    self.net.append(nn.Linear(__UpperCAmelCase						,					__UpperCAmelCase					)					)
    # FF as used in Vision Transformer, MLP-Mixer, etc. have a final dropout
    if final_dropout:
       self.net.append(nn.Dropout(__UpperCAmelCase					)					)




 def 							lowerCamelCase   (    self						,					__UpperCAmelCase					):



    '''simple docstring'''
    for module in self.net:
       __lowerCamelCase    				=     module(__UpperCAmelCase					)
    return hidden_states



class 			__lowerCAmelCase  (						nn.Module ):




 def __init__(    self						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase = "none"					):



    '''simple docstring'''
    super().__init__()
    __lowerCamelCase    				=     nn.Linear(__UpperCAmelCase						,					__UpperCAmelCase					)
    __lowerCamelCase    				=     approximate




 def 							lowerCamelCase   (    self						,					__UpperCAmelCase					):



    '''simple docstring'''
    if gate.device.type != "mps":
       return F.gelu(__UpperCAmelCase						,					approximate=self.approximate					)
    # mps: gelu is not implemented for float16
    return F.gelu(gate.to(dtype=torch.floataa					)						,					approximate=self.approximate					).to(dtype=gate.dtype					)




 def 							lowerCamelCase   (    self						,					__UpperCAmelCase					):



    '''simple docstring'''
    __lowerCamelCase    				=     self.proj(__UpperCAmelCase					)
    __lowerCamelCase    				=     self.gelu(__UpperCAmelCase					)
    return hidden_states



class 			__lowerCAmelCase  (						nn.Module ):




 def __init__(    self						,					__UpperCAmelCase						,					__UpperCAmelCase					):



    '''simple docstring'''
    super().__init__()
    __lowerCamelCase    				=     nn.Linear(__UpperCAmelCase						,					dim_out * 2					)




 def 							lowerCamelCase   (    self						,					__UpperCAmelCase					):



    '''simple docstring'''
    if gate.device.type != "mps":
       return F.gelu(__UpperCAmelCase					)
    # mps: gelu is not implemented for float16
    return F.gelu(gate.to(dtype=torch.floataa					)					).to(dtype=gate.dtype					)




 def 							lowerCamelCase   (    self						,					__UpperCAmelCase					):



    '''simple docstring'''
    __lowerCamelCase				,__lowerCamelCase    				=     self.proj(__UpperCAmelCase					).chunk(2						,					dim=-1					)
    return hidden_states * self.gelu(__UpperCAmelCase					)



class 			__lowerCAmelCase  (						nn.Module ):




 def __init__(    self						,					__UpperCAmelCase						,					__UpperCAmelCase					):



    '''simple docstring'''
    super().__init__()
    __lowerCamelCase    				=     nn.Linear(__UpperCAmelCase						,					__UpperCAmelCase					)




 def 							lowerCamelCase   (    self						,					__UpperCAmelCase					):



    '''simple docstring'''
    __lowerCamelCase    				=     self.proj(__UpperCAmelCase					)
    return x * torch.sigmoid(1.702 * x					)



class 			__lowerCAmelCase  (						nn.Module ):




 def __init__(    self						,					__UpperCAmelCase						,					__UpperCAmelCase					):



    '''simple docstring'''
    super().__init__()
    __lowerCamelCase    				=     nn.Embedding(__UpperCAmelCase						,					__UpperCAmelCase					)
    __lowerCamelCase    				=     nn.SiLU()
    __lowerCamelCase    				=     nn.Linear(__UpperCAmelCase						,					embedding_dim * 2					)
    __lowerCamelCase    				=     nn.LayerNorm(__UpperCAmelCase						,					elementwise_affine=__UpperCAmelCase					)




 def 							lowerCamelCase   (    self						,					__UpperCAmelCase						,					__UpperCAmelCase					):



    '''simple docstring'''
    __lowerCamelCase    				=     self.linear(self.silu(self.emb(__UpperCAmelCase					)					)					)
    __lowerCamelCase				,__lowerCamelCase    				=     torch.chunk(__UpperCAmelCase						,					2					)
    __lowerCamelCase    				=     self.norm(__UpperCAmelCase					) * (1 + scale) + shift
    return x



class 			__lowerCAmelCase  (						nn.Module ):




 def __init__(    self						,					__UpperCAmelCase						,					__UpperCAmelCase					):



    '''simple docstring'''
    super().__init__()

    __lowerCamelCase    				=     CombinedTimestepLabelEmbeddings(__UpperCAmelCase						,					__UpperCAmelCase					)

    __lowerCamelCase    				=     nn.SiLU()
    __lowerCamelCase    				=     nn.Linear(__UpperCAmelCase						,					6 * embedding_dim						,					bias=__UpperCAmelCase					)
    __lowerCamelCase    				=     nn.LayerNorm(__UpperCAmelCase						,					elementwise_affine=__UpperCAmelCase						,					eps=1E-6					)




 def 							lowerCamelCase   (    self						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase=None					):



    '''simple docstring'''
    __lowerCamelCase    				=     self.linear(self.silu(self.emb(__UpperCAmelCase						,					__UpperCAmelCase						,					hidden_dtype=__UpperCAmelCase					)					)					)
    __lowerCamelCase				,__lowerCamelCase				,__lowerCamelCase				,__lowerCamelCase				,__lowerCamelCase				,__lowerCamelCase    				=     emb.chunk(6						,					dim=1					)
    __lowerCamelCase    				=     self.norm(__UpperCAmelCase					) * (1 + scale_msa[:, None]) + shift_msa[:, None]
    return x, gate_msa, shift_mlp, scale_mlp, gate_mlp



class 			__lowerCAmelCase  (						nn.Module ):




 def __init__(    self						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase = None						,					__UpperCAmelCase = 1E-5					):



    '''simple docstring'''
    super().__init__()
    __lowerCamelCase    				=     num_groups
    __lowerCamelCase    				=     eps

    if act_fn is None:
       __lowerCamelCase    				=     None
    else:
       __lowerCamelCase    				=     get_activation(__UpperCAmelCase					)

    __lowerCamelCase    				=     nn.Linear(__UpperCAmelCase						,					out_dim * 2					)




 def 							lowerCamelCase   (    self						,					__UpperCAmelCase						,					__UpperCAmelCase					):



    '''simple docstring'''
    if self.act:
       __lowerCamelCase    				=     self.act(__UpperCAmelCase					)
    __lowerCamelCase    				=     self.linear(__UpperCAmelCase					)
    __lowerCamelCase    				=     emb[:, :, None, None]
    __lowerCamelCase				,__lowerCamelCase    				=     emb.chunk(2						,					dim=1					)

    __lowerCamelCase    				=     F.group_norm(__UpperCAmelCase						,					self.num_groups						,					eps=self.eps					)
    __lowerCamelCase    				=     x * (1 + scale) + shift
    return x











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


a_								=       {
    """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[Any]				):
   __lowerCamelCase    				=     {}
   state_dict.pop('''pixel_mean'''			,_UpperCamelCase				)
   state_dict.pop('''pixel_std'''			,_UpperCamelCase				)

   __lowerCamelCase    				=     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    				=     key.replace(_UpperCamelCase			,_UpperCamelCase				)

      if re.match(_UpperCamelCase			,_UpperCamelCase				):
         __lowerCamelCase    				=     int(re.match(_UpperCamelCase			,_UpperCamelCase				).group(2				)				)
         if layer_nb == 0:
            __lowerCamelCase    				=     key.replace('''layers.0'''			,'''proj_in'''				)
         elif layer_nb == 1:
            __lowerCamelCase    				=     key.replace('''layers.1'''			,'''layers.0'''				)
         elif layer_nb == 2:
            __lowerCamelCase    				=     key.replace('''layers.2'''			,'''proj_out'''				)

      __lowerCamelCase    				=     value

   __lowerCamelCase    				=     model_state_dict[
       '''prompt_encoder.shared_embedding.positional_embedding'''
   ]

   return model_state_dict





def 		a__      (		_UpperCamelCase      :			Tuple			,_UpperCamelCase      :			Dict			,_UpperCamelCase      :			Any			,_UpperCamelCase      :			Optional[int]="ybelkada/segment-anything"				):
   __lowerCamelCase    				=     hf_hub_download(_UpperCamelCase			,F"""checkpoints/{model_name}.pth"""				)

   if "sam_vit_b" in model_name:
      __lowerCamelCase    				=     SamConfig()
   elif "sam_vit_l" in model_name:
      __lowerCamelCase    				=     SamVisionConfig(
          hidden_size=10_24			,num_hidden_layers=24			,num_attention_heads=16			,global_attn_indexes=[5, 11, 17, 23]			,)

      __lowerCamelCase    				=     SamConfig(
          vision_config=_UpperCamelCase			,)
   elif "sam_vit_h" in model_name:
      __lowerCamelCase    				=     SamVisionConfig(
          hidden_size=12_80			,num_hidden_layers=32			,num_attention_heads=16			,global_attn_indexes=[7, 15, 23, 31]			,)

      __lowerCamelCase    				=     SamConfig(
          vision_config=_UpperCamelCase			,)

   __lowerCamelCase    				=     torch.load(_UpperCamelCase			,map_location='''cpu'''				)
   __lowerCamelCase    				=     replace_keys(_UpperCamelCase				)

   __lowerCamelCase    				=     SamImageProcessor()

   __lowerCamelCase    				=     SamProcessor(image_processor=_UpperCamelCase				)
   __lowerCamelCase    				=     SamModel(_UpperCamelCase				)

   hf_model.load_state_dict(_UpperCamelCase				)
   __lowerCamelCase    				=     hf_model.to('''cuda'''				)

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

   __lowerCamelCase    				=     [[[4_00, 6_50]]]
   __lowerCamelCase    				=     [[1]]

   __lowerCamelCase    				=     processor(images=np.array(_UpperCamelCase				)			,return_tensors='''pt'''				).to('''cuda'''				)

   with torch.no_grad():
      __lowerCamelCase    				=     hf_model(**_UpperCamelCase				)
   __lowerCamelCase    				=     output.iou_scores.squeeze()

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

      __lowerCamelCase    				=     processor(
          images=np.array(_UpperCamelCase				)			,input_points=_UpperCamelCase			,input_labels=_UpperCamelCase			,return_tensors='''pt'''				).to('''cuda'''				)

      with torch.no_grad():
         __lowerCamelCase    				=     hf_model(**_UpperCamelCase				)
      __lowerCamelCase    				=     output.iou_scores.squeeze()

      assert scores[-1].item() == 0.9_712_603_092_193_604

      __lowerCamelCase    				=     ((75, 2_75, 17_25, 8_50),)

      __lowerCamelCase    				=     processor(images=np.array(_UpperCamelCase				)			,input_boxes=_UpperCamelCase			,return_tensors='''pt'''				).to('''cuda'''				)

      with torch.no_grad():
         __lowerCamelCase    				=     hf_model(**_UpperCamelCase				)
      __lowerCamelCase    				=     output.iou_scores.squeeze()

      assert scores[-1].item() == 0.8_686_015_605_926_514

      # Test with 2 points and 1 image.
      __lowerCamelCase    				=     [[[4_00, 6_50], [8_00, 6_50]]]
      __lowerCamelCase    				=     [[1, 1]]

      __lowerCamelCase    				=     processor(
          images=np.array(_UpperCamelCase				)			,input_points=_UpperCamelCase			,input_labels=_UpperCamelCase			,return_tensors='''pt'''				).to('''cuda'''				)

      with torch.no_grad():
         __lowerCamelCase    				=     hf_model(**_UpperCamelCase				)
      __lowerCamelCase    				=     output.iou_scores.squeeze()

      assert scores[-1].item() == 0.9_936_047_792_434_692


if __name__ == "__main__":
  a_								=       argparse.ArgumentParser()
  a_								=       ["""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""",
  )

  a_								=       parser.parse_args()

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






