import glob
import os
from typing import Dict, List, Union

import torch

from diffusers.utils import is_safetensors_available


if is_safetensors_available():
 import safetensors.torch

from huggingface_hub import snapshot_download

from diffusers import DiffusionPipeline, __version__
from diffusers.schedulers.scheduling_utils import SCHEDULER_CONFIG_NAME
from diffusers.utils import CONFIG_NAME, DIFFUSERS_CACHE, ONNX_WEIGHTS_NAME, WEIGHTS_NAME




class  __lowerCAmelCase	(							a_    ):



 '''simple docstring'''

 def __init__(      self							:		List[Any]	)			->				Optional[int]:
        self.register_to_config()
        super().__init__()

 def        __snake_case       (      self							:		List[Any]							,      lowerCamelCase							:		Any							,      lowerCamelCase							:		Any	)			->				List[str]:
        if dicta == dicta:
               return True
        else:
               _UpperCamelCase    ,   _UpperCamelCase  :							Tuple	    =	self._remove_meta_keys(lowerCamelCase	)
               _UpperCamelCase    ,   _UpperCamelCase  :							Any	    =	self._remove_meta_keys(lowerCamelCase	)
               if configa == configa:
                      print(f"Warning !: Mismatch in keys {meta_keysa} and {meta_keysa}."	)
                      return True
        return False

 def        __snake_case       (      self							:		Any							,      lowerCamelCase							:		Dict	)			->				str:
        _UpperCamelCase  :							List[str]	    =	[]
        _UpperCamelCase  :							Optional[Any]	    =	config_dict.copy()
        for key in config_dict.keys():
               if key.startswith('_'	):
                      temp_dict.pop(lowerCamelCase	)
                      meta_keys.append(lowerCamelCase	)
        return (temp_dict, meta_keys)

 @torch.no_grad()
 def        __snake_case       (      self							:		str							,      lowerCamelCase							:		List[Union[str, os.PathLike]]							,      **lowerCamelCase							:		Optional[Any]	)			->				Any:
        _UpperCamelCase  :							List[Any]	    =	kwargs.pop('cache_dir'							,      lowerCamelCase	)
        _UpperCamelCase  :							Any	    =	kwargs.pop('resume_download'							,      lowerCamelCase	)
        _UpperCamelCase  :							Optional[int]	    =	kwargs.pop('force_download'							,      lowerCamelCase	)
        _UpperCamelCase  :							Union[str, Any]	    =	kwargs.pop('proxies'							,      lowerCamelCase	)
        _UpperCamelCase  :							Union[str, Any]	    =	kwargs.pop('local_files_only'							,      lowerCamelCase	)
        _UpperCamelCase  :							Union[str, Any]	    =	kwargs.pop('use_auth_token'							,      lowerCamelCase	)
        _UpperCamelCase  :							Optional[int]	    =	kwargs.pop('revision'							,      lowerCamelCase	)
        _UpperCamelCase  :							Any	    =	kwargs.pop('torch_dtype'							,      lowerCamelCase	)
        _UpperCamelCase  :							Any	    =	kwargs.pop('device_map'							,      lowerCamelCase	)

        _UpperCamelCase  :							List[str]	    =	kwargs.pop('alpha'							,      0.5	)
        _UpperCamelCase  :							List[Any]	    =	kwargs.pop('interp'							,      lowerCamelCase	)

        print('Received list'							,      lowerCamelCase	)
        print(f"Combining with alpha={alpha}, interpolation mode={interp}"	)

        _UpperCamelCase  :							Any	    =	len(lowerCamelCase	)
        # Ignore result from model_index_json comparision of the two checkpoints
        _UpperCamelCase  :							Dict	    =	kwargs.pop('force'							,      lowerCamelCase	)

        # If less than 2 checkpoints, nothing to merge. If more than 3, not supported for now.
        if checkpoint_count > 3 or checkpoint_count < 2:
               raise ValueError(
                   'Received incorrect number of checkpoints to merge. Ensure that either 2 or 3 checkpoints are being'
                   ' passed.'	)

        print('Received the right number of checkpoints'	)
        # chkpt0, chkpt1 = pretrained_model_name_or_path_list[0:2]
        # chkpt2 = pretrained_model_name_or_path_list[2] if checkpoint_count == 3 else None

        # Validate that the checkpoints can be merged
        # Step 1: Load the model config and compare the checkpoints. We'll compare the model_index.json first while ignoring the keys starting with '_'
        _UpperCamelCase  :							List[Any]	    =	[]
        for pretrained_model_name_or_path in pretrained_model_name_or_path_list:
               _UpperCamelCase  :							Tuple	    =	DiffusionPipeline.load_config(
                   lowerCamelCase							,      cache_dir=lowerCamelCase							,      resume_download=lowerCamelCase							,      force_download=lowerCamelCase							,      proxies=lowerCamelCase							,      local_files_only=lowerCamelCase							,      use_auth_token=lowerCamelCase							,      revision=lowerCamelCase							,      )
               config_dicts.append(lowerCamelCase	)

        _UpperCamelCase  :							Tuple	    =	True
        for idx in range(1							,      len(lowerCamelCase	)	):
               comparison_result &= self._compare_model_configs(config_dicts[idx - 1]							,      config_dicts[idx]	)
               if not force and comparison_result is False:
                      raise ValueError('Incompatible checkpoints. Please check model_index.json for the models.'	)
                      print(config_dicts[0]							,      config_dicts[1]	)
        print('Compatible model_index.json files found'	)
        # Step 2: Basic Validation has succeeded. Let's download the models and save them into our local files.
        _UpperCamelCase  :							Union[str, Any]	    =	[]
        for pretrained_model_name_or_path, config_dict in zip(lowerCamelCase							,      lowerCamelCase	):
               _UpperCamelCase  :							str	    =	[k for k in config_dict.keys() if not k.startswith('_'	)]
               _UpperCamelCase  :							List[Any]	    =	[os.path.join(lowerCamelCase							,      '*'	) for k in folder_names]
               allow_patterns += [
                   WEIGHTS_NAME,
                   SCHEDULER_CONFIG_NAME,
                   CONFIG_NAME,
                   ONNX_WEIGHTS_NAME,
                   DiffusionPipeline.config_name,
               ]
               _UpperCamelCase  :							Dict	    =	config_dict.get('_class_name'	)
               _UpperCamelCase  :							Optional[int]	    =	{'diffusers': __version__, 'pipeline_class': requested_pipeline_class}

               _UpperCamelCase  :							Any	    =	(
                   pretrained_model_name_or_path
                   if os.path.isdir(lowerCamelCase	)
                   else snapshot_download(
                       lowerCamelCase							,      cache_dir=lowerCamelCase							,      resume_download=lowerCamelCase							,      proxies=lowerCamelCase							,      local_files_only=lowerCamelCase							,      use_auth_token=lowerCamelCase							,      revision=lowerCamelCase							,      allow_patterns=lowerCamelCase							,      user_agent=lowerCamelCase							,      )
               )
               print('Cached Folder'							,      lowerCamelCase	)
               cached_folders.append(lowerCamelCase	)

        # Step 3:-
        # Load the first checkpoint as a diffusion pipeline and modify its module state_dict in place
        _UpperCamelCase  :							int	    =	DiffusionPipeline.from_pretrained(
            cached_folders[0]							,      torch_dtype=lowerCamelCase							,      device_map=lowerCamelCase	)
        final_pipe.to(self.device	)

        _UpperCamelCase  :							str	    =	None
        if len(lowerCamelCase	) > 2:
               _UpperCamelCase  :							Union[str, Any]	    =	os.path.join(cached_folders[2]	)

        if interp == "sigmoid":
               _UpperCamelCase  :							Any	    =	CheckpointMergerPipeline.sigmoid
        elif interp == "inv_sigmoid":
               _UpperCamelCase  :							Tuple	    =	CheckpointMergerPipeline.inv_sigmoid
        elif interp == "add_diff":
               _UpperCamelCase  :							List[Any]	    =	CheckpointMergerPipeline.add_difference
        else:
               _UpperCamelCase  :							Tuple	    =	CheckpointMergerPipeline.weighted_sum

        # Find each module's state dict.
        for attr in final_pipe.config.keys():
               if not attr.startswith('_'	):
                      _UpperCamelCase  :							Dict	    =	os.path.join(cached_folders[1]							,      lowerCamelCase	)
                      if os.path.exists(lowerCamelCase	):
                             _UpperCamelCase  :							Any	    =	[
                                 *glob.glob(os.path.join(lowerCamelCase							,      '*.safetensors'	)	),
                                 *glob.glob(os.path.join(lowerCamelCase							,      '*.bin'	)	),
                             ]
                             _UpperCamelCase  :							Optional[int]	    =	files[0] if len(lowerCamelCase	) > 0 else None
                      if len(lowerCamelCase	) < 3:
                             _UpperCamelCase  :							Tuple	    =	None
                      else:
                             _UpperCamelCase  :							int	    =	os.path.join(cached_folders[2]							,      lowerCamelCase	)
                             if os.path.exists(lowerCamelCase	):
                                    _UpperCamelCase  :							Dict	    =	[
                                        *glob.glob(os.path.join(lowerCamelCase							,      '*.safetensors'	)	),
                                        *glob.glob(os.path.join(lowerCamelCase							,      '*.bin'	)	),
                                    ]
                                    _UpperCamelCase  :							Optional[int]	    =	files[0] if len(lowerCamelCase	) > 0 else None
                # For an attr if both checkpoint_path_1 and 2 are None, ignore.
                # If atleast one is present, deal with it according to interp method, of course only if the state_dict keys match.
                      if checkpoint_path_a is None and checkpoint_path_a is None:
                             print(f"Skipping {attr}: not present in 2nd or 3d model"	)
                             continue
                      try:
                             _UpperCamelCase  :							Tuple	    =	getattr(lowerCamelCase							,      lowerCamelCase	)
                             if isinstance(lowerCamelCase							,      lowerCamelCase	):  # ignore requires_safety_checker boolean
                                    continue
                             _UpperCamelCase  :							Optional[int]	    =	getattr(lowerCamelCase							,      'state_dict'	)
                             _UpperCamelCase  :							Union[str, Any]	    =	theta_a()

                             _UpperCamelCase  :							Optional[int]	    =	getattr(lowerCamelCase							,      'load_state_dict'	)
                             _UpperCamelCase  :							Dict	    =	(
                                 safetensors.torch.load_file(lowerCamelCase	)
                                 if (is_safetensors_available() and checkpoint_path_a.endswith('.safetensors'	))
                                 else torch.load(lowerCamelCase							,      map_location='cpu'	)
                             )
                             _UpperCamelCase  :							List[str]	    =	None
                             if checkpoint_path_a:
                                    _UpperCamelCase  :							int	    =	(
                                        safetensors.torch.load_file(lowerCamelCase	)
                                        if (is_safetensors_available() and checkpoint_path_a.endswith('.safetensors'	))
                                        else torch.load(lowerCamelCase							,      map_location='cpu'	)
                                    )

                             if not theta_a.keys() == theta_a.keys():
                                    print(f"Skipping {attr}: key mismatch"	)
                                    continue
                             if theta_a and not theta_a.keys() == theta_a.keys():
                                    print(f"Skipping {attr}:y mismatch"	)
                      except Exception as e:
                             print(f"Skipping {attr} do to an unexpected error: {str(lowerCamelCase	)}"	)
                             continue
                      print(f"MERGING {attr}"	)

                      for key in theta_a.keys():
                             if theta_a:
                                    _UpperCamelCase  :							Union[str, Any]	    =	theta_func(theta_a[key]							,      theta_a[key]							,      theta_a[key]							,      lowerCamelCase	)
                             else:
                                    _UpperCamelCase  :							Dict	    =	theta_func(theta_a[key]							,      theta_a[key]							,      lowerCamelCase							,      lowerCamelCase	)

                      del theta_a
                      del theta_a
                      update_theta_a(lowerCamelCase	)

                      del theta_a
        return final_pipe

 @staticmethod
 def        __snake_case       (      lowerCamelCase							:		Optional[int]							,      lowerCamelCase							:		Dict							,      lowerCamelCase							:		str							,      lowerCamelCase							:		Optional[Any]	)			->				Dict:
        return ((1 - alpha) * thetaa) + (alpha * thetaa)

 @staticmethod
 def        __snake_case       (      lowerCamelCase							:		Any							,      lowerCamelCase							:		Union[str, Any]							,      lowerCamelCase							:		Union[str, Any]							,      lowerCamelCase							:		Dict	)			->				str:
        _UpperCamelCase  :							Union[str, Any]	    =	alpha * alpha * (3 - (2 * alpha))
        return thetaa + ((thetaa - thetaa) * alpha)

 @staticmethod
 def        __snake_case       (      lowerCamelCase							:		Union[str, Any]							,      lowerCamelCase							:		Optional[Any]							,      lowerCamelCase							:		int							,      lowerCamelCase							:		int	)			->				List[str]:
        import math

        _UpperCamelCase  :							str	    =	0.5 - math.sin(math.asin(1.0 - 2.0 * alpha	) / 3.0	)
        return thetaa + ((thetaa - thetaa) * alpha)

 @staticmethod
 def        __snake_case       (      lowerCamelCase							:		int							,      lowerCamelCase							:		Any							,      lowerCamelCase							:		str							,      lowerCamelCase							:		Tuple	)			->				int:
        return thetaa + (thetaa - thetaa) * (1.0 - alpha)




import argparse

import torch

from transformers import ChineseCLIPConfig, ChineseCLIPModel


def SCREAMING_SNAKE_CASE					(  lowercase_:						List[str]		,    lowercase_:						Dict		,    lowercase_:						List[Any]					):
       _UpperCamelCase    ,   _UpperCamelCase    ,   _UpperCamelCase  :							List[Any]	    =	pt_weights[F"{prefix}.in_proj_weight"].chunk(3		,    dim=0					)
       _UpperCamelCase    ,   _UpperCamelCase    ,   _UpperCamelCase  :							Dict	    =	pt_weights[F"{prefix}.in_proj_bias"].chunk(3		,    dim=0					)

       _UpperCamelCase  :							List[str]	    =	pt_weights[F"{prefix}.out_proj.weight"]
       _UpperCamelCase  :							Any	    =	pt_weights[F"{prefix}.out_proj.bias"]

       _UpperCamelCase  :							List[Any]	    =	q_proj
       _UpperCamelCase  :							Union[str, Any]	    =	q_proj_bias

       _UpperCamelCase  :							str	    =	k_proj
       _UpperCamelCase  :							List[str]	    =	k_proj_bias

       _UpperCamelCase  :							Dict	    =	v_proj
       _UpperCamelCase  :							str	    =	v_proj_bias

       _UpperCamelCase  :							List[str]	    =	out_proj_weights
       _UpperCamelCase  :							Any	    =	out_proj_bias


def SCREAMING_SNAKE_CASE					(  lowercase_:						List[str]		,    lowercase_:						List[Any]		,    lowercase_:						Any					):
       copy_linear(hf_mlp.fca		,    lowercase_		,    F"{prefix}.c_fc"					)
       copy_linear(hf_mlp.fca		,    lowercase_		,    F"{prefix}.c_proj"					)


def SCREAMING_SNAKE_CASE					(  lowercase_:						Optional[Any]		,    lowercase_:						List[str]		,    lowercase_:						Any					):
       _UpperCamelCase  :							Dict	    =	pt_weights[F"{prefix}.weight"].data
       _UpperCamelCase  :							List[Any]	    =	pt_weights[F"{prefix}.bias"].data


def SCREAMING_SNAKE_CASE					(  lowercase_:						Any		,    lowercase_:						Optional[int]		,    lowercase_:						Union[str, Any]					):
       # copy layer norms
       copy_linear(hf_layer.layer_norma		,    lowercase_		,    F"{prefix}.ln_1"					)
       copy_linear(hf_layer.layer_norma		,    lowercase_		,    F"{prefix}.ln_2"					)

       # copy MLP
       copy_mlp(hf_layer.mlp		,    lowercase_		,    F"{prefix}.mlp"					)

       # copy attn
       copy_attn_layer(hf_layer.self_attn		,    lowercase_		,    F"{prefix}.attn"					)


def SCREAMING_SNAKE_CASE					(  lowercase_:						str		,    lowercase_:						str		,    lowercase_:						int					):
       for layer_id, hf_layer in enumerate(lowercase_					):
              copy_layer(lowercase_		,    lowercase_		,    F"{prefix}.{layer_id}"					)


def SCREAMING_SNAKE_CASE					(  lowercase_:						Union[str, Any]		,    lowercase_:						List[Any]					):
       # copy projection
       _UpperCamelCase  :							List[str]	    =	pt_weights['text_projection'].data.T

       # copy text encoder
       for name, param in hf_model.text_model.named_parameters():
              _UpperCamelCase  :							Tuple	    =	pt_weights[F"bert.{name}"].data


def SCREAMING_SNAKE_CASE					(  lowercase_:						str		,    lowercase_:						Any					):
       # copy projection
       _UpperCamelCase  :							List[Any]	    =	pt_weights['visual.proj'].data.T

       # copy layer norms
       copy_linear(hf_model.vision_model.pre_layrnorm		,    lowercase_		,    'visual.ln_pre'					)
       copy_linear(hf_model.vision_model.post_layernorm		,    lowercase_		,    'visual.ln_post'					)

       # copy embeddings
       _UpperCamelCase  :							List[Any]	    =	pt_weights['visual.conv1.weight'].data
       _UpperCamelCase  :							List[Any]	    =	pt_weights['visual.class_embedding'].data
       _UpperCamelCase  :							int	    =	pt_weights['visual.positional_embedding'].data

       # copy encoder
       copy_layers(hf_model.vision_model.encoder.layers		,    lowercase_		,    'visual.transformer.resblocks'					)


@torch.no_grad()
def SCREAMING_SNAKE_CASE					(  lowercase_:						Optional[Any]		,    lowercase_:						Optional[int]		,    lowercase_:						List[Any]=None					):
       assert config_path is not None, "Please specify the ChineseCLIP model config of the corresponding model size."
       _UpperCamelCase  :							Optional[Any]	    =	ChineseCLIPConfig.from_pretrained(lowercase_					)

       _UpperCamelCase  :							List[Any]	    =	ChineseCLIPModel(lowercase_					).eval()

       _UpperCamelCase  :							List[Any]	    =	torch.load(lowercase_		,    map_location='cpu'					)['state_dict']
       _UpperCamelCase  :							Union[str, Any]	    =	{(name[7:] if name.startswith('module.'					) else name): value for name, value in pt_weights.items()}

       copy_text_model_and_projection(lowercase_		,    lowercase_					)
       copy_vision_model_and_projection(lowercase_		,    lowercase_					)
       _UpperCamelCase  :							List[Any]	    =	pt_weights['logit_scale'].data

       hf_model.save_pretrained(lowercase_					)


if __name__ == "__main__":
 UpperCAmelCase      		=					argparse.ArgumentParser()
 parser.add_argument(
     '''--pytorch_dump_folder_path''',
     default=None,
     type=str,
     help='''Path to the output folder storing converted hf PyTorch model.''',
 )
 parser.add_argument(
     '''--checkpoint_path''', default=None, type=str, help='''Path to original github format ChineseCLIP checkpoint.'''
 )
 parser.add_argument(
     '''--config_path''', default=None, required=True, type=str, help='''Path to hf config.json of model to convert.'''
 )
 UpperCAmelCase      		=					parser.parse_args()

 convert_chinese_clip_checkpoint(args.checkpoint_path, args.pytorch_dump_folder_path, args.config_path)
 print('''The conversion is finished!''')



