



import random
import unittest

import numpy as np
import torch

from diffusers import (
    DPMSolverMultistepScheduler,
    EulerAncestralDiscreteScheduler,
    EulerDiscreteScheduler,
    LMSDiscreteScheduler,
    OnnxStableDiffusionUpscalePipeline,
    PNDMScheduler,
)
from diffusers.utils import floats_tensor
from diffusers.utils.testing_utils import (
    is_onnx_available,
    load_image,
    nightly,
    require_onnxruntime,
    require_torch_gpu,
)

from ..test_pipelines_onnx_common import OnnxPipelineTesterMixin


if is_onnx_available():
  import onnxruntime as ort



class 			__lowerCAmelCase  (						lowerCAmelCase__	,						unittest.TestCase ):
 # TODO: is there an appropriate internal test set?
 lowerCAmelCase__							       =	"""ssube/stable-diffusion-x4-upscaler-onnx"""




 def 							lowerCamelCase   (    self						,					__UpperCAmelCase=0					):



    '''simple docstring'''
    __lowerCamelCase    				=     floats_tensor((1, 3, 128, 128)						,					rng=random.Random(__UpperCAmelCase					)					)
    __lowerCamelCase    				=     torch.manual_seed(__UpperCAmelCase					)
    __lowerCamelCase    				=     {
        '''prompt''': '''A painting of a squirrel eating a burger''',
        '''image''': image,
        '''generator''': generator,
        '''num_inference_steps''': 3,
        '''guidance_scale''': 7.5,
        '''output_type''': '''numpy''',
    }
    return inputs




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    __lowerCamelCase    				=     OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint						,					provider='''CPUExecutionProvider'''					)
    pipe.set_progress_bar_config(disable=__UpperCAmelCase					)

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

    # started as 128, should now be 512
    assert image.shape == (1, 512, 512, 3)
    __lowerCamelCase    				=     np.array(
        [0.6_974_782, 0.68_902_093, 0.70_135_885, 0.7_583_618, 0.7_804_545, 0.7_854_912, 0.78_667_426, 0.78_743_863, 0.78_070_223]					)
    assert np.abs(image_slice - expected_slice					).max() < 1E-1




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    __lowerCamelCase    				=     OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint						,					provider='''CPUExecutionProvider'''					)
    __lowerCamelCase    				=     PNDMScheduler.from_config(pipe.scheduler.config						,					skip_prk_steps=__UpperCAmelCase					)
    pipe.set_progress_bar_config(disable=__UpperCAmelCase					)

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

    assert image.shape == (1, 512, 512, 3)
    __lowerCamelCase    				=     np.array(
        [0.6_898_892, 0.59_240_556, 0.52_499_527, 0.58_866_215, 0.52_258_235, 0.52_572_715, 0.62_414_473, 0.6_174_387, 0.6_214_964]					)
    assert np.abs(image_slice.flatten() - expected_slice					).max() < 1E-1




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    __lowerCamelCase    				=     OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint						,					provider='''CPUExecutionProvider'''					)
    __lowerCamelCase    				=     DPMSolverMultistepScheduler.from_config(pipe.scheduler.config					)
    pipe.set_progress_bar_config(disable=__UpperCAmelCase					)

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

    assert image.shape == (1, 512, 512, 3)
    __lowerCamelCase    				=     np.array(
        [0.7_659_278, 0.76_437_664, 0.75_579_107, 0.7_691_116, 0.77_666_986, 0.7_727_672, 0.7_758_664, 0.7_812_226, 0.76_942_515]					)

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




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    __lowerCamelCase    				=     OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint						,					provider='''CPUExecutionProvider'''					)
    __lowerCamelCase    				=     EulerDiscreteScheduler.from_config(pipe.scheduler.config					)
    pipe.set_progress_bar_config(disable=__UpperCAmelCase					)

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

    assert image.shape == (1, 512, 512, 3)
    __lowerCamelCase    				=     np.array(
        [0.6_974_782, 0.68_902_093, 0.70_135_885, 0.7_583_618, 0.7_804_545, 0.7_854_912, 0.78_667_426, 0.78_743_863, 0.78_070_223]					)
    assert np.abs(image_slice.flatten() - expected_slice					).max() < 1E-1




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    __lowerCamelCase    				=     OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint						,					provider='''CPUExecutionProvider'''					)
    __lowerCamelCase    				=     EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config					)
    pipe.set_progress_bar_config(disable=__UpperCAmelCase					)

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

    assert image.shape == (1, 512, 512, 3)
    __lowerCamelCase    				=     np.array(
        [0.77_424_496, 0.773_601, 0.7_645_288, 0.7_769_598, 0.7_772_739, 0.7_738_688, 0.78_187_233, 0.77_879_584, 0.767_043]					)

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



@nightly
@require_onnxruntime
@require_torch_gpu
class 			__lowerCAmelCase  (						unittest.TestCase ):




 @property
 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    return (
        "CUDAExecutionProvider",
        {
            "gpu_mem_limit": "15000000000",  # 15GB
            "arena_extend_strategy": "kSameAsRequested",
        },
    )




 @property
 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    __lowerCamelCase    				=     ort.SessionOptions()
    __lowerCamelCase    				=     False
    return options




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    __lowerCamelCase    				=     load_image(
        '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
        '''/img2img/sketch-mountains-input.jpg'''					)
    __lowerCamelCase    				=     init_image.resize((128, 128)					)
    # using the PNDM scheduler by default
    __lowerCamelCase    				=     OnnxStableDiffusionUpscalePipeline.from_pretrained(
        '''ssube/stable-diffusion-x4-upscaler-onnx'''						,					provider=self.gpu_provider						,					sess_options=self.gpu_options						,					)
    pipe.set_progress_bar_config(disable=__UpperCAmelCase					)

    __lowerCamelCase    				=     '''A fantasy landscape, trending on artstation'''

    __lowerCamelCase    				=     torch.manual_seed(0					)
    __lowerCamelCase    				=     pipe(
        prompt=__UpperCAmelCase						,					image=__UpperCAmelCase						,					guidance_scale=7.5						,					num_inference_steps=10						,					generator=__UpperCAmelCase						,					output_type='''np'''						,					)
    __lowerCamelCase    				=     output.images
    __lowerCamelCase    				=     images[0, 255:258, 383:386, -1]

    assert images.shape == (1, 512, 512, 3)
    __lowerCamelCase    				=     np.array([0.4_883, 0.4_947, 0.4_980, 0.4_975, 0.4_982, 0.4_980, 0.5_000, 0.5_006, 0.4_972]					)
    # TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    __lowerCamelCase    				=     load_image(
        '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
        '''/img2img/sketch-mountains-input.jpg'''					)
    __lowerCamelCase    				=     init_image.resize((128, 128)					)
    __lowerCamelCase    				=     LMSDiscreteScheduler.from_pretrained(
        '''ssube/stable-diffusion-x4-upscaler-onnx'''						,					subfolder='''scheduler'''					)
    __lowerCamelCase    				=     OnnxStableDiffusionUpscalePipeline.from_pretrained(
        '''ssube/stable-diffusion-x4-upscaler-onnx'''						,					scheduler=__UpperCAmelCase						,					provider=self.gpu_provider						,					sess_options=self.gpu_options						,					)
    pipe.set_progress_bar_config(disable=__UpperCAmelCase					)

    __lowerCamelCase    				=     '''A fantasy landscape, trending on artstation'''

    __lowerCamelCase    				=     torch.manual_seed(0					)
    __lowerCamelCase    				=     pipe(
        prompt=__UpperCAmelCase						,					image=__UpperCAmelCase						,					guidance_scale=7.5						,					num_inference_steps=20						,					generator=__UpperCAmelCase						,					output_type='''np'''						,					)
    __lowerCamelCase    				=     output.images
    __lowerCamelCase    				=     images[0, 255:258, 383:386, -1]

    assert images.shape == (1, 512, 512, 3)
    __lowerCamelCase    				=     np.array(
        [0.50_173_753, 0.50_223_356, 0.502_039, 0.50_233_036, 0.5_023_725, 0.5_022_601, 0.5_018_758, 0.50_234_085, 0.50_241_566]					)
    # TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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











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)






