




'''simple docstring'''



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   ( __lowerCAmelCase):


       def __init__(   self  )      ->  Optional[Any]:



            '''simple docstring'''




            self.register_to_config()
            super().__init__()


       def       lowerCAmelCase      (   self    ,  __SCREAMING_SNAKE_CASE    ,  __SCREAMING_SNAKE_CASE  )      ->  Any:



            '''simple docstring'''




            if dicta == dicta:
                 return True
            else:
                 __snake_case    					=  self._remove_meta_keys(__lowerCamelCase  )
                 __snake_case    					=  self._remove_meta_keys(__lowerCamelCase  )
                 if configa == configa:
                      print(F'''Warning !: Mismatch in keys {meta_keysa} and {meta_keysa}.'''  )
                      return True
            return False


       def       lowerCAmelCase      (   self    ,  __SCREAMING_SNAKE_CASE  )      ->  Tuple:



            '''simple docstring'''




            __snake_case    					=  []
            __snake_case    					=  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       lowerCAmelCase      (   self    ,  __SCREAMING_SNAKE_CASE    ,  **__SCREAMING_SNAKE_CASE  )      ->  List[Any]:



            '''simple docstring'''




            __snake_case    					=  kwargs.pop('''cache_dir'''    ,  __lowerCamelCase  )
            __snake_case    					=  kwargs.pop('''resume_download'''    ,  __lowerCamelCase  )
            __snake_case    					=  kwargs.pop('''force_download'''    ,  __lowerCamelCase  )
            __snake_case    					=  kwargs.pop('''proxies'''    ,  __lowerCamelCase  )
            __snake_case    					=  kwargs.pop('''local_files_only'''    ,  __lowerCamelCase  )
            __snake_case    					=  kwargs.pop('''use_auth_token'''    ,  __lowerCamelCase  )
            __snake_case    					=  kwargs.pop('''revision'''    ,  __lowerCamelCase  )
            __snake_case    					=  kwargs.pop('''torch_dtype'''    ,  __lowerCamelCase  )
            __snake_case    					=  kwargs.pop('''device_map'''    ,  __lowerCamelCase  )

            __snake_case    					=  kwargs.pop('''alpha'''    ,  0.5  )
            __snake_case    					=  kwargs.pop('''interp'''    ,  __lowerCamelCase  )

            print('''Received list'''    ,  __lowerCamelCase  )
            print(F'''Combining with alpha={alpha}, interpolation mode={interp}'''  )

            __snake_case    					=  len(__lowerCamelCase  )
            # Ignore result from model_index_json comparision of the two checkpoints
            __snake_case    					=  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 '_'
            __snake_case    					=  []
            for pretrained_model_name_or_path in pretrained_model_name_or_path_list:
                 __snake_case    					=  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  )

            __snake_case    					=  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.
            __snake_case    					=  []
            for pretrained_model_name_or_path, config_dict in zip(__lowerCamelCase    ,  __lowerCamelCase  ):
                 __snake_case    					=  [k for k in config_dict.keys() if not k.startswith('''_'''  )]
                 __snake_case    					=  [os.path.join(__lowerCamelCase    ,  '''*'''  ) for k in folder_names]
                 allow_patterns += [
                     WEIGHTS_NAME,
                     SCHEDULER_CONFIG_NAME,
                     CONFIG_NAME,
                     ONNX_WEIGHTS_NAME,
                     DiffusionPipeline.config_name,
                 ]
                 __snake_case    					=  config_dict.get('''_class_name'''  )
                 __snake_case    					=  {"diffusers": __version__, "pipeline_class": requested_pipeline_class}

                 __snake_case    					=  (
                     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
            __snake_case    					=  DiffusionPipeline.from_pretrained(
                cached_folders[0]    ,  torch_dtype=__lowerCamelCase    ,  device_map=__lowerCamelCase  )
            final_pipe.to(self.device  )

            __snake_case    					=  None
            if len(__lowerCamelCase  ) > 2:
                 __snake_case    					=  os.path.join(cached_folders[2]  )

            if interp == "sigmoid":
                 __snake_case    					=  CheckpointMergerPipeline.sigmoid
            elif interp == "inv_sigmoid":
                 __snake_case    					=  CheckpointMergerPipeline.inv_sigmoid
            elif interp == "add_diff":
                 __snake_case    					=  CheckpointMergerPipeline.add_difference
            else:
                 __snake_case    					=  CheckpointMergerPipeline.weighted_sum

            # Find each module's state dict.
            for attr in final_pipe.config.keys():
                 if not attr.startswith('''_'''  ):
                      __snake_case    					=  os.path.join(cached_folders[1]    ,  __lowerCamelCase  )
                      if os.path.exists(__lowerCamelCase  ):
                           __snake_case    					=  [
                               *glob.glob(os.path.join(__lowerCamelCase    ,  '''*.safetensors'''  )  ),
                               *glob.glob(os.path.join(__lowerCamelCase    ,  '''*.bin'''  )  ),
                           ]
                           __snake_case    					=  files[0] if len(__lowerCamelCase  ) > 0 else None
                      if len(__lowerCamelCase  ) < 3:
                           __snake_case    					=  None
                      else:
                           __snake_case    					=  os.path.join(cached_folders[2]    ,  __lowerCamelCase  )
                           if os.path.exists(__lowerCamelCase  ):
                                __snake_case    					=  [
                                    *glob.glob(os.path.join(__lowerCamelCase    ,  '''*.safetensors'''  )  ),
                                    *glob.glob(os.path.join(__lowerCamelCase    ,  '''*.bin'''  )  ),
                                ]
                                __snake_case    					=  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:
                           __snake_case    					=  getattr(__lowerCamelCase    ,  __lowerCamelCase  )
                           if isinstance(__lowerCamelCase    ,  __lowerCamelCase  ):  # ignore requires_safety_checker boolean
                                continue
                           __snake_case    					=  getattr(__lowerCamelCase    ,  '''state_dict'''  )
                           __snake_case    					=  theta_a()

                           __snake_case    					=  getattr(__lowerCamelCase    ,  '''load_state_dict'''  )
                           __snake_case    					=  (
                               safetensors.torch.load_file(__lowerCamelCase  )
                               if (is_safetensors_available() and checkpoint_path_a.endswith('''.safetensors'''  ))
                               else torch.load(__lowerCamelCase    ,  map_location='''cpu'''  )
                           )
                           __snake_case    					=  None
                           if checkpoint_path_a:
                                __snake_case    					=  (
                                    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:
                                __snake_case    					=  theta_func(theta_a[key]    ,  theta_a[key]    ,  theta_a[key]    ,  __lowerCamelCase  )
                           else:
                                __snake_case    					=  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       lowerCAmelCase      (   __SCREAMING_SNAKE_CASE    ,  __SCREAMING_SNAKE_CASE    ,  __SCREAMING_SNAKE_CASE    ,  __SCREAMING_SNAKE_CASE  )      ->  str:



            '''simple docstring'''




            return ((1 - alpha) * thetaa) + (alpha * thetaa)


       @staticmethod
       def       lowerCAmelCase      (   __SCREAMING_SNAKE_CASE    ,  __SCREAMING_SNAKE_CASE    ,  __SCREAMING_SNAKE_CASE    ,  __SCREAMING_SNAKE_CASE  )      ->  List[Any]:



            '''simple docstring'''




            __snake_case    					=  alpha * alpha * (3 - (2 * alpha))
            return thetaa + ((thetaa - thetaa) * alpha)


       @staticmethod
       def       lowerCAmelCase      (   __SCREAMING_SNAKE_CASE    ,  __SCREAMING_SNAKE_CASE    ,  __SCREAMING_SNAKE_CASE    ,  __SCREAMING_SNAKE_CASE  )      ->  Union[str, Any]:



            '''simple docstring'''




            import math

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




       @staticmethod
       def       lowerCAmelCase      (   __SCREAMING_SNAKE_CASE    ,  __SCREAMING_SNAKE_CASE    ,  __SCREAMING_SNAKE_CASE    ,  __SCREAMING_SNAKE_CASE  )      ->  Optional[Any]:



            '''simple docstring'''




            return thetaa + (thetaa - thetaa) * (1.0 - alpha)








'''simple docstring'''



# Copyright 2022 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import os
from typing import List







def 			_UpperCamelCase (_lowerCamelCase							:     List[str]   ,	_lowerCamelCase							:     str	)->   str:




     '''simple docstring'''





     if name != "training_function" and name != "main":
          raise ValueError(f'''Incorrect function name passed: {name}, choose either \'main\' or \'training_function\''''	)
     __snake_case					,    __snake_case    					=  [], False
     for line in lines:
          if not found_start and f'''def {name}''' in line:
               __snake_case    					=  True
               good_lines.append(_lowerCamelCase	)
               continue
          if found_start:
               if name == "training_function" and "def main" in line:
                    return good_lines
               if name == "main" and "if __name__" in line:
                    return good_lines
               good_lines.append(_lowerCamelCase	)







def 			_UpperCamelCase (_lowerCamelCase							:     List[str]	)->   int:




     '''simple docstring'''





     return [line for line in lines if not line.lstrip().startswith('''#'''	) and line != "\n"]







def 			_UpperCamelCase (_lowerCamelCase							:     str   ,	_lowerCamelCase							:     str   ,	_lowerCamelCase							:     bool   ,	_lowerCamelCase							:     str = None	)->   Optional[Any]:




     '''simple docstring'''





     with open(_lowerCamelCase   ,	'''r'''	) as f:
          __snake_case    					=  f.readlines()
     with open(os.path.abspath(os.path.join('''examples'''   ,	'''nlp_example.py'''	)	)   ,	'''r'''	) as f:
          __snake_case    					=  f.readlines()
     with open(_lowerCamelCase   ,	'''r'''	) as f:
          __snake_case    					=  f.readlines()
     if secondary_filename is not None:
          with open(_lowerCamelCase   ,	'''r'''	) as f:
               __snake_case    					=  f.readlines()

    # This is our base, we remove all the code from here in our `full_filename` and `feature_filename` to find the new content
     if parser_only:
          __snake_case    					=  clean_lines(get_function_contents_by_name(_lowerCamelCase   ,	'''main'''	)	)
          __snake_case    					=  clean_lines(get_function_contents_by_name(_lowerCamelCase   ,	'''main'''	)	)
          __snake_case    					=  clean_lines(get_function_contents_by_name(_lowerCamelCase   ,	'''main'''	)	)
          if secondary_filename is not None:
               __snake_case    					=  clean_lines(get_function_contents_by_name(_lowerCamelCase   ,	'''main'''	)	)
     else:
          __snake_case    					=  clean_lines(get_function_contents_by_name(_lowerCamelCase   ,	'''training_function'''	)	)
          __snake_case    					=  clean_lines(get_function_contents_by_name(_lowerCamelCase   ,	'''training_function'''	)	)
          __snake_case    					=  clean_lines(get_function_contents_by_name(_lowerCamelCase   ,	'''training_function'''	)	)
          if secondary_filename is not None:
               __snake_case    					=  clean_lines(
                   get_function_contents_by_name(_lowerCamelCase   ,	'''training_function'''	)	)

     __snake_case    					=  '''train_dataloader, eval_dataloader = get_dataloaders(accelerator, batch_size)\n'''

     # Specific code in our script that differs from the full version, aka what is new
     __snake_case    					=  []
     __snake_case    					=  []  # We keep track of the idxs just in case it's a repeated statement
     __snake_case    					=  iter(_lowerCamelCase	)
     for i in range(len(_lowerCamelCase	) - 1	):
          if i not in passed_idxs:
               __snake_case    					=  next(_lowerCamelCase	)
               if (line not in full_file_func) and (line.lstrip() != _dl_line):
                    if "TESTING_MOCKED_DATALOADERS" not in line:
                         new_feature_code.append(_lowerCamelCase	)
                         passed_idxs.append(_lowerCamelCase	)
                    else:
                         # Skip over the `config['num_epochs'] = 2` statement
                         __snake_case    					=  next(_lowerCamelCase	)

    # Extract out just the new parts from the full_file_training_func
     __snake_case    					=  []
     __snake_case    					=  []  # We keep track of the idxs just in case it's a repeated statement
     for i, line in enumerate(_lowerCamelCase	):
          if i not in passed_idxs:
               if (line not in full_file_func) and (line.lstrip() != _dl_line):
                    if "TESTING_MOCKED_DATALOADERS" not in line:
                         new_full_example_parts.append(_lowerCamelCase	)
                         passed_idxs.append(_lowerCamelCase	)

    # Finally, get the overall diff
     __snake_case    					=  [line for line in new_feature_code if line not in new_full_example_parts]
     if secondary_filename is not None:
          __snake_case    					=  [line for line in full_file_contents if line not in secondary_file_func]
          __snake_case    					=  [line for line in diff_from_example if line not in diff_from_two]

     return diff_from_example






