




'''simple docstring'''


import argparse

import requests
import torch
from PIL import Image

from transformers import (
    CLIPSegConfig,
    CLIPSegForImageSegmentation,
    CLIPSegProcessor,
    CLIPSegTextConfig,
    CLIPSegVisionConfig,
    CLIPTokenizer,
    ViTImageProcessor,
)
def 					SCREAMING_SNAKE_CASE(   __lowercase	)			-> List[str]:
     A:		Optional[int]          =		CLIPSegTextConfig()
     A:		int          =		CLIPSegVisionConfig(patch_size=1_6	)

     A:		Union[str, Any]          =		True if '''refined''' in model_name else False
     A:		int          =		1_6 if '''rd16''' in model_name else 6_4

     A:		str          =		CLIPSegConfig.from_text_vision_configs(
         __lowercase		,				__lowercase		,				use_complex_transposed_convolution=__lowercase		,				reduce_dim=__lowercase		,				)
     return config
def 					SCREAMING_SNAKE_CASE(   __lowercase	)			-> Dict:
     # update prefixes
     if "clip_model" in name:
          A:		Any          =		name.replace('''clip_model'''		,				'''clip'''	)
     if "transformer" in name:
          if "visual" in name:
               A:		int          =		name.replace('''visual.transformer'''		,				'''vision_model'''	)
          else:
               A:		Dict          =		name.replace('''transformer'''		,				'''text_model'''	)
     if "resblocks" in name:
          A:		Tuple          =		name.replace('''resblocks'''		,				'''encoder.layers'''	)
     if "ln_1" in name:
          A:		int          =		name.replace('''ln_1'''		,				'''layer_norm1'''	)
     if "ln_2" in name:
          A:		Union[str, Any]          =		name.replace('''ln_2'''		,				'''layer_norm2'''	)
     if "c_fc" in name:
          A:		Tuple          =		name.replace('''c_fc'''		,				'''fc1'''	)
     if "c_proj" in name:
          A:		Union[str, Any]          =		name.replace('''c_proj'''		,				'''fc2'''	)
     if "attn" in name and "self" not in name:
          A:		str          =		name.replace('''attn'''		,				'''self_attn'''	)
     # text encoder
     if "token_embedding" in name:
          A:		Optional[int]          =		name.replace('''token_embedding'''		,				'''text_model.embeddings.token_embedding'''	)
     if "positional_embedding" in name and "visual" not in name:
          A:		int          =		name.replace('''positional_embedding'''		,				'''text_model.embeddings.position_embedding.weight'''	)
     if "ln_final" in name:
          A:		str          =		name.replace('''ln_final'''		,				'''text_model.final_layer_norm'''	)
     # vision encoder
     if "visual.class_embedding" in name:
          A:		Optional[Any]          =		name.replace('''visual.class_embedding'''		,				'''vision_model.embeddings.class_embedding'''	)
     if "visual.conv1" in name:
          A:		int          =		name.replace('''visual.conv1'''		,				'''vision_model.embeddings.patch_embedding'''	)
     if "visual.positional_embedding" in name:
          A:		Optional[Any]          =		name.replace('''visual.positional_embedding'''		,				'''vision_model.embeddings.position_embedding.weight'''	)
     if "visual.ln_pre" in name:
          A:		str          =		name.replace('''visual.ln_pre'''		,				'''vision_model.pre_layrnorm'''	)
     if "visual.ln_post" in name:
          A:		Dict          =		name.replace('''visual.ln_post'''		,				'''vision_model.post_layernorm'''	)
     # projection layers
     if "visual.proj" in name:
          A:		Any          =		name.replace('''visual.proj'''		,				'''visual_projection.weight'''	)
     if "text_projection" in name:
          A:		int          =		name.replace('''text_projection'''		,				'''text_projection.weight'''	)
     # decoder
     if "trans_conv" in name:
          A:		Tuple          =		name.replace('''trans_conv'''		,				'''transposed_convolution'''	)
     if "film_mul" in name or "film_add" in name or "reduce" in name or "transposed_convolution" in name:
          A:		Optional[Any]          =		'''decoder.''' + name
     if "blocks" in name:
          A:		Any          =		name.replace('''blocks'''		,				'''decoder.layers'''	)
     if "linear1" in name:
          A:		Optional[int]          =		name.replace('''linear1'''		,				'''mlp.fc1'''	)
     if "linear2" in name:
          A:		List[str]          =		name.replace('''linear2'''		,				'''mlp.fc2'''	)
     if "norm1" in name and "layer_" not in name:
          A:		Optional[int]          =		name.replace('''norm1'''		,				'''layer_norm1'''	)
     if "norm2" in name and "layer_" not in name:
          A:		Any          =		name.replace('''norm2'''		,				'''layer_norm2'''	)

     return name
def 					SCREAMING_SNAKE_CASE(   __lowercase		,				__lowercase	)			-> Optional[int]:
     for key in orig_state_dict.copy().keys():
          A:		List[str]          =		orig_state_dict.pop(__lowercase	)

          if key.startswith('''clip_model'''	) and "attn.in_proj" in key:
               A:		Any          =		key.split('''.'''	)
               if "visual" in key:
                    A:		Optional[int]          =		int(key_split[4]	)
                    A:		List[str]          =		config.vision_config.hidden_size
                    A:		Optional[int]          =		'''vision_model'''
               else:
                    A:		List[str]          =		int(key_split[3]	)
                    A:		str          =		config.text_config.hidden_size
                    A:		int          =		'''text_model'''

               if "weight" in key:
                    A:		List[str]          =		val[:dim, :]
                    A:		Optional[int]          =		val[
                        dim : dim * 2, :
                    ]
                    A:		Any          =		val[-dim:, :]
               else:
                    A:		Any          =		val[:dim]
                    A:		Any          =		val[dim : dim * 2]
                    A:		Any          =		val[-dim:]
          elif "self_attn" in key and "out_proj" not in key:
               A:		Any          =		key.split('''.'''	)
               A:		List[Any]          =		int(key_split[1]	)
               A:		List[Any]          =		config.reduce_dim
               if "weight" in key:
                    A:		List[Any]          =		val[:dim, :]
                    A:		List[Any]          =		val[dim : dim * 2, :]
                    A:		int          =		val[-dim:, :]
               else:
                    A:		Optional[int]          =		val[:dim]
                    A:		int          =		val[dim : dim * 2]
                    A:		int          =		val[-dim:]
          else:
               A:		List[str]          =		rename_key(__lowercase	)
               if "visual_projection" in new_name or "text_projection" in new_name:
                    A:		Dict          =		val.T
               A:		Tuple          =		val

     return orig_state_dict
def 					SCREAMING_SNAKE_CASE(   )			-> Optional[int]:
     A:		int          =		'''http://images.cocodataset.org/val2017/000000039769.jpg'''
     A:		Any          =		Image.open(requests.get(__lowercase		,				stream=__lowercase	).raw	)
     return image
def 					SCREAMING_SNAKE_CASE(   __lowercase		,				__lowercase		,				__lowercase		,				__lowercase	)			-> Union[str, Any]:
     A:		Optional[Any]          =		get_clipseg_config(__lowercase	)
     A:		List[str]          =		CLIPSegForImageSegmentation(__lowercase	)
     model.eval()

     A:		Any          =		torch.load(__lowercase		,				map_location='''cpu'''	)

     # remove some keys
     for key in state_dict.copy().keys():
          if key.startswith('''model'''	):
               state_dict.pop(__lowercase		,				__lowercase	)

    # rename some keys
     A:		Tuple          =		convert_state_dict(__lowercase		,				__lowercase	)
     A     ,    A:		Optional[int]          =		model.load_state_dict(__lowercase		,				strict=__lowercase	)

     if missing_keys != ["clip.text_model.embeddings.position_ids", "clip.vision_model.embeddings.position_ids"]:
          raise ValueError('''Missing keys that are not expected: {}'''.format(__lowercase	)	)
     if unexpected_keys != ["decoder.reduce.weight", "decoder.reduce.bias"]:
          raise ValueError(F"""Unexpected keys: {unexpected_keys}"""	)

     A:		Tuple          =		ViTImageProcessor(size=3_5_2	)
     A:		Dict          =		CLIPTokenizer.from_pretrained('''openai/clip-vit-base-patch32'''	)
     A:		Tuple          =		CLIPSegProcessor(image_processor=__lowercase		,				tokenizer=__lowercase	)

     A:		Optional[int]          =		prepare_img()
     A:		str          =		['''a glass''', '''something to fill''', '''wood''', '''a jar''']

     A:		Dict          =		processor(text=__lowercase		,				images=[image] * len(__lowercase	)		,				padding='''max_length'''		,				return_tensors='''pt'''	)

     with torch.no_grad():
          A:		Any          =		model(**__lowercase	)

     # verify values
     A:		List[Any]          =		torch.tensor([0.1_1_1_0, -0.1_8_8_2, 0.1_6_4_5]	)
     A:		Union[str, Any]          =		torch.tensor([0.2_6_9_2, -0.7_1_9_7, -0.1_3_2_8]	)
     if model_name == "clipseg-rd64-refined":
          A:		Optional[int]          =		torch.tensor(
              [[-1_0.0_4_0_7, -9.9_4_3_1, -1_0.2_6_4_6], [-9.9_7_5_1, -9.7_0_6_4, -9.9_5_8_6], [-9.6_8_9_1, -9.5_6_4_5, -9.9_6_1_8]]	)
     elif model_name == "clipseg-rd64":
          A:		Any          =		torch.tensor(
              [[-7.2_8_7_7, -7.2_7_1_1, -7.2_4_6_3], [-7.2_6_5_2, -7.2_7_8_0, -7.2_5_2_0], [-7.2_2_3_9, -7.2_2_0_4, -7.2_0_0_1]]	)
     elif model_name == "clipseg-rd16":
          A:		List[Any]          =		torch.tensor(
              [[-6.3_9_5_5, -6.4_0_5_5, -6.4_1_5_1], [-6.3_9_1_1, -6.4_0_3_3, -6.4_1_0_0], [-6.3_4_7_4, -6.3_7_0_2, -6.3_7_6_2]]	)
     else:
          raise ValueError(F"""Model name {model_name} not supported."""	)

     assert torch.allclose(outputs.logits[0, :3, :3]		,				__lowercase		,				atol=1E-3	)
     assert torch.allclose(outputs.conditional_embeddings[0, :3]		,				__lowercase		,				atol=1E-3	)
     assert torch.allclose(outputs.pooled_output[0, :3]		,				__lowercase		,				atol=1E-3	)
     print('''Looks ok!'''	)

     if pytorch_dump_folder_path is not None:
          print(F"""Saving model and processor to {pytorch_dump_folder_path}"""	)
          model.save_pretrained(__lowercase	)
          processor.save_pretrained(__lowercase	)

     if push_to_hub:
          print(F"""Pushing model and processor for {model_name} to the hub"""	)
          model.push_to_hub(F"""CIDAS/{model_name}"""	)
          processor.push_to_hub(F"""CIDAS/{model_name}"""	)


if __name__ == "__main__":
     UpperCamelCase			     =					argparse.ArgumentParser()
     # Required parameters
     parser.add_argument(
         '''--model_name''',
         default='''clipseg-rd64''',
         type=str,
         choices=['''clipseg-rd16''', '''clipseg-rd64''', '''clipseg-rd64-refined'''],
         help=(
             '''Name of the model. Supported models are: clipseg-rd64, clipseg-rd16 and clipseg-rd64-refined (rd meaning'''
             ''' reduce dimension)'''
         ),
     )
     parser.add_argument(
         '''--checkpoint_path''',
         default='''/Users/nielsrogge/Documents/CLIPSeg/clip_plus_rd64-uni.pth''',
         type=str,
         help=(
             '''Path to the original checkpoint. Note that the script assumes that the checkpoint includes both CLIP and'''
             ''' the decoder weights.'''
         ),
     )
     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.'''
     )

     UpperCamelCase			     =					parser.parse_args()
     convert_clipseg_checkpoint(args.model_name, args.checkpoint_path, args.pytorch_dump_folder_path, args.push_to_hub)












'''simple docstring'''


import timeit
from typing import Callable, Optional

from ..configuration_utils import PretrainedConfig
from ..models.auto.modeling_auto import MODEL_MAPPING, MODEL_WITH_LM_HEAD_MAPPING
from ..utils import is_pyanvml_available, is_torch_available, logging
from .benchmark_utils import (
    Benchmark,
    Memory,
    MemorySummary,
    measure_peak_memory_cpu,
    start_memory_tracing,
    stop_memory_tracing,
)


if is_torch_available():
     import torch

     from .benchmark_args import PyTorchBenchmarkArguments


if is_pyanvml_available():
     import pyanvml.pyanvml as nvml


UpperCamelCase			     =					logging.get_logger(__name__)






class 						lowerCAmelCase_	(       UpperCAmelCase_							):







  '''simple docstring'''

  UpperCamelCase_			:      PyTorchBenchmarkArguments
  UpperCamelCase_			:      PretrainedConfig
  UpperCamelCase_			:      str										=				"PyTorch"



  @property
  def 					_snake_case     ( self				:							int	)			->		List[str]:







       '''simple docstring'''



       return torch.__version__



  def 					_snake_case     ( self				:							Dict					,  SCREAMING_SNAKE_CASE_				:							str					,  SCREAMING_SNAKE_CASE_				:							int					,  SCREAMING_SNAKE_CASE_				:							int	)			->		float:







       '''simple docstring'''



       A:		List[str]          =		self._prepare_inference_func(SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_	)
       return self._measure_speed(_inference	)



  def 					_snake_case     ( self				:							Dict					,  SCREAMING_SNAKE_CASE_				:							str					,  SCREAMING_SNAKE_CASE_				:							int					,  SCREAMING_SNAKE_CASE_				:							int	)			->		[Memory, Optional[MemorySummary]]:







       '''simple docstring'''



       A:		Any          =		self._prepare_inference_func(SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_	)
       return self._measure_memory(_inference	)



  def 					_snake_case     ( self				:							List[str]					,  SCREAMING_SNAKE_CASE_				:							str					,  SCREAMING_SNAKE_CASE_				:							int					,  SCREAMING_SNAKE_CASE_				:							int	)			->		float:







       '''simple docstring'''



       A:		List[Any]          =		self._prepare_train_func(SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_	)
       return self._measure_speed(_train	)



  def 					_snake_case     ( self				:							int					,  SCREAMING_SNAKE_CASE_				:							str					,  SCREAMING_SNAKE_CASE_				:							int					,  SCREAMING_SNAKE_CASE_				:							int	)			->		[Memory, Optional[MemorySummary]]:







       '''simple docstring'''



       A:		List[Any]          =		self._prepare_train_func(SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_	)
       return self._measure_memory(_train	)



  def 					_snake_case     ( self				:							Union[str, Any]					,  SCREAMING_SNAKE_CASE_				:							str					,  SCREAMING_SNAKE_CASE_				:							int					,  SCREAMING_SNAKE_CASE_				:							int	)			->		Callable[[], None]:







       '''simple docstring'''



       A:		int          =		self.config_dict[model_name]

       if self.args.torchscript:
            A:		int          =		True

       A:		Any          =		(
           hasattr(SCREAMING_SNAKE_CASE_					,  '''architectures'''	)
           and isinstance(config.architectures					,  SCREAMING_SNAKE_CASE_	)
           and len(config.architectures	) > 0
       )
       if not self.args.only_pretrain_model and has_model_class_in_config:
            try:
                 A:		Optional[int]          =		config.architectures[0]
                 A:		Optional[int]          =		__import__('''transformers'''					,  fromlist=[model_class]	)
                 A:		Dict          =		getattr(SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_	)
                 A:		Optional[Any]          =		model_cls(SCREAMING_SNAKE_CASE_	)
            except ImportError:
                 raise ImportError(
                     f"""{model_class} does not exist. If you just want to test the pretrained model, you might want to"""
                     ''' set `--only_pretrain_model` or `args.only_pretrain_model=True`.'''	)
       else:
            A:		Union[str, Any]          =		MODEL_MAPPING[config.__class__](SCREAMING_SNAKE_CASE_	)

       model.eval()
       model.to(self.args.device	)

       # encoder-decoder has vocab size saved differently
       A:		Tuple          =		config.vocab_size if hasattr(SCREAMING_SNAKE_CASE_					,  '''vocab_size'''	) else config.encoder.vocab_size
       A:		List[Any]          =		torch.randint(SCREAMING_SNAKE_CASE_					,  (batch_size, sequence_length)					,  dtype=torch.long					,  device=self.args.device	)

       if self.args.fpaa:
            logger.info('''Running training in Mixed Precision...'''	)
            if not self.args.is_gpu:
                 raise ValueError('''Mixed precision is possible only for GPU.'''	)
            # amp seems to have memory leaks so that memory usage
            # is measured using .half() for now https://github.com/NVIDIA/apex/issues/439
            model.half()

       if self.args.torchscript:
            with torch.no_grad():
                 A:		Any          =		torch.jit.trace(SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_	)
       else:
            A:		Dict          =		model

       def encoder_decoder_forward():
            with torch.no_grad():
                 A:		List[str]          =		inference_model(SCREAMING_SNAKE_CASE_					,  decoder_input_ids=SCREAMING_SNAKE_CASE_	)
            return outputs

       def encoder_forward():
            with torch.no_grad():
                 A:		Union[str, Any]          =		inference_model(SCREAMING_SNAKE_CASE_	)
            return outputs

       A:		List[Any]          =		encoder_decoder_forward if config.is_encoder_decoder else encoder_forward
       return _forward



  def 					_snake_case     ( self				:							Tuple					,  SCREAMING_SNAKE_CASE_				:							str					,  SCREAMING_SNAKE_CASE_				:							int					,  SCREAMING_SNAKE_CASE_				:							int	)			->		Callable[[], None]:







       '''simple docstring'''



       A:		int          =		self.config_dict[model_name]

       A:		str          =		(
           hasattr(SCREAMING_SNAKE_CASE_					,  '''architectures'''	)
           and isinstance(config.architectures					,  SCREAMING_SNAKE_CASE_	)
           and len(config.architectures	) > 0
       )
       if not self.args.only_pretrain_model and has_model_class_in_config:
            try:
                 A:		List[str]          =		config.architectures[0]
                 A:		List[str]          =		__import__('''transformers'''					,  fromlist=[model_class]	)
                 A:		Optional[int]          =		getattr(SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_	)
                 A:		Union[str, Any]          =		model_cls(SCREAMING_SNAKE_CASE_	)
            except ImportError:
                 raise ImportError(
                     f"""{model_class} does not exist. If you just want to test the pretrained model, you might want to"""
                     ''' set `--only_pretrain_model` or `args.only_pretrain_model=True`.'''	)
       else:
            A:		Union[str, Any]          =		MODEL_WITH_LM_HEAD_MAPPING[config.__class__](SCREAMING_SNAKE_CASE_	)

       if self.args.torchscript:
            raise NotImplementedError('''Training for torchscript is currently not implemented'''	)
       else:
            A:		str          =		model

       model.train()
       model.to(self.args.device	)

       # encoder-decoder has vocab size saved differently
       A:		Tuple          =		config.vocab_size if hasattr(SCREAMING_SNAKE_CASE_					,  '''vocab_size'''	) else config.encoder.vocab_size
       A:		str          =		torch.randint(SCREAMING_SNAKE_CASE_					,  (batch_size, sequence_length)					,  dtype=torch.long					,  device=self.args.device	)

       if self.args.fpaa:
            logger.info('''Running training in Mixed Precision...'''	)
            if not self.args.is_gpu:
                 raise ValueError('''Mixed precision is possible only for GPU.'''	)

            # amp seems to have memory leaks so that memory usage
            # is measured using .half() for now https://github.com/NVIDIA/apex/issues/439
            model.half()

       def compute_loss_and_backprob_encoder():
            A:		str          =		train_model(SCREAMING_SNAKE_CASE_					,  labels=SCREAMING_SNAKE_CASE_	)[0]
            loss.backward()
            return loss

       def compute_loss_and_backprob_encoder_decoder():
            A:		Dict          =		train_model(SCREAMING_SNAKE_CASE_					,  decoder_input_ids=SCREAMING_SNAKE_CASE_					,  labels=SCREAMING_SNAKE_CASE_	)[0]
            loss.backward()
            return loss

       A:		str          =		(
           compute_loss_and_backprob_encoder_decoder
           if config.is_encoder_decoder
           else compute_loss_and_backprob_encoder
       )
       return _train



  def 					_snake_case     ( self				:							Any					,  SCREAMING_SNAKE_CASE_				:							Any	)			->		float:







       '''simple docstring'''



       try:
            if self.args.is_tpu or self.args.torchscript:
                 # run additional 10 times to stabilize compilation for tpu and torchscript
                 logger.info('''Do inference on TPU or torchscript. Running model 5 times to stabilize compilation'''	)
                 timeit.repeat(
                     SCREAMING_SNAKE_CASE_					,  repeat=1					,  number=5					,  )

            # as written in https://docs.python.org/2/library/timeit.html#timeit.Timer.repeat, min should be taken rather than the average
            A:		Any          =		timeit.repeat(
                SCREAMING_SNAKE_CASE_					,  repeat=self.args.repeat					,  number=10					,  )

            if self.args.is_tpu and self.args.torch_xla_tpu_print_metrics:
                 import torch_xla.debug.metrics as met

                 self.print_fn(met.metrics_report()	)

            return min(SCREAMING_SNAKE_CASE_	) / 10.0
       except RuntimeError as e:
            self.print_fn(f"""Doesn't fit on GPU. {e}"""	)
            return "N/A"




  def 					_snake_case     ( self				:							Optional[int]					,  SCREAMING_SNAKE_CASE_				:							Callable[[], None]	)			->		[Memory, MemorySummary]:







       '''simple docstring'''



       try:
            if self.args.trace_memory_line_by_line:
                 A:		List[Any]          =		start_memory_tracing('''transformers'''	)

            if self.args.is_tpu:
                 # tpu
                 raise NotImplementedError(
                     '''Memory Benchmarking is currently not implemented for TPU. Please disable memory benchmarking with'''
                     ''' `--no-memory` or `args.memory=False`'''	)
            elif self.args.is_gpu:
                 if not is_pyanvml_available():
                      logger.warning(
                          '''py3nvml not installed, we won\'t log GPU memory usage. '''
                          '''Install py3nvml (pip install py3nvml) to log information about GPU.'''	)
                      A:		Optional[int]          =		'''N/A'''
                 else:
                      logger.info(
                          '''Measuring total GPU usage on GPU device. Make sure to not have additional processes running'''
                          ''' on the same GPU.'''	)
                      # init nvml
                      nvml.nvmlInit()
                      func()
                      A:		Union[str, Any]          =		nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx	)
                      A:		Union[str, Any]          =		nvml.nvmlDeviceGetMemoryInfo(SCREAMING_SNAKE_CASE_	)
                      A:		Tuple          =		meminfo.used
                      A:		Tuple          =		Memory(SCREAMING_SNAKE_CASE_	)
                      # shutdown nvml
                      nvml.nvmlShutdown()
            else:
                 # cpu
                 A:		int          =		measure_peak_memory_cpu(SCREAMING_SNAKE_CASE_	)
                 A:		str          =		Memory(SCREAMING_SNAKE_CASE_	) if isinstance(SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_	) else memory_bytes

            if self.args.trace_memory_line_by_line:
                 A:		Dict          =		stop_memory_tracing(SCREAMING_SNAKE_CASE_	)
            else:
                 A:		List[str]          =		None

            return memory, summary
       except RuntimeError as e:
            self.print_fn(f"""Doesn't fit on GPU. {e}"""	)
            return "N/A", None






