"""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 					UpperCamelCase   (							lowerCAmelCase__  ):







       def __init__(					self) -> List[str]:
            self.register_to_config()
            super().__init__()







       def 						a_		(					self,				lowerCAmelCase__,				lowerCAmelCase__) -> int:
            if dicta == dicta:
                 return True
            else:
                 snake_case_			,      snake_case_   			=   self._remove_meta_keys(lowerCAmelCase__)
                 snake_case_			,      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 						a_		(					self,				lowerCAmelCase__) -> List[Any]:
            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 						a_		(					self,				lowerCAmelCase__,				**lowerCAmelCase__) -> Any:
            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 						a_		(					lowerCAmelCase__,				lowerCAmelCase__,				lowerCAmelCase__,				lowerCAmelCase__) -> Optional[Any]:
            return ((1 - alpha) * thetaa) + (alpha * thetaa)







       @staticmethod
       def 						a_		(					lowerCAmelCase__,				lowerCAmelCase__,				lowerCAmelCase__,				lowerCAmelCase__) -> Any:
            snake_case_   			=   alpha * alpha * (3 - (2 * alpha))
            return thetaa + ((thetaa - thetaa) * alpha)







       @staticmethod
       def 						a_		(					lowerCAmelCase__,				lowerCAmelCase__,				lowerCAmelCase__,				lowerCAmelCase__) -> Union[str, Any]:
            import math

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







       @staticmethod
       def 						a_		(					lowerCAmelCase__,				lowerCAmelCase__,				lowerCAmelCase__,				lowerCAmelCase__) -> List[str]:
            return thetaa + (thetaa - thetaa) * (1.0 - alpha)


"""simple docstring"""


import posixpath
from pathlib import Path
from unittest.mock import patch

import pytest
from fsspec.implementations.local import AbstractFileSystem, LocalFileSystem, stringify_path
from fsspec.registry import _registry as _fsspec_registry




class 					UpperCamelCase   (							lowerCAmelCase__  ):
       SCREAMING_SNAKE_CASE_						      = "mock"







       def __init__(					self,				*lowerCAmelCase__,				lowerCAmelCase__,				**lowerCAmelCase__) -> Dict:
            super().__init__()
            snake_case_   			=   LocalFileSystem(*lowerCAmelCase__,				**lowerCAmelCase__)
            snake_case_   			=   Path(lowerCAmelCase__).resolve().as_posix() + '/'







       def 						a_		(					self,				lowerCAmelCase__,				*lowerCAmelCase__,				**lowerCAmelCase__) -> Any:
            snake_case_   			=   posixpath.join(self.local_root_dir,				self._strip_protocol(lowerCAmelCase__))
            return self._fs.mkdir(lowerCAmelCase__,				*lowerCAmelCase__,				**lowerCAmelCase__)







       def 						a_		(					self,				lowerCAmelCase__,				*lowerCAmelCase__,				**lowerCAmelCase__) -> str:
            snake_case_   			=   posixpath.join(self.local_root_dir,				self._strip_protocol(lowerCAmelCase__))
            return self._fs.makedirs(lowerCAmelCase__,				*lowerCAmelCase__,				**lowerCAmelCase__)







       def 						a_		(					self,				lowerCAmelCase__) -> Union[str, Any]:
            snake_case_   			=   posixpath.join(self.local_root_dir,				self._strip_protocol(lowerCAmelCase__))
            return self._fs.rmdir(lowerCAmelCase__)







       def 						a_		(					self,				lowerCAmelCase__,				lowerCAmelCase__=True,				*lowerCAmelCase__,				**lowerCAmelCase__) -> Optional[Any]:
            snake_case_   			=   posixpath.join(self.local_root_dir,				self._strip_protocol(lowerCAmelCase__))
            snake_case_   			=   self._fs.ls(lowerCAmelCase__,				detail=lowerCAmelCase__,				*lowerCAmelCase__,				**lowerCAmelCase__)
            if detail:
                 return [{**info, "name": info["name"][len(self.local_root_dir) :]} for info in out]
            else:
                 return [name[len(self.local_root_dir) :] for name in out]







       def 						a_		(					self,				lowerCAmelCase__,				*lowerCAmelCase__,				**lowerCAmelCase__) -> Dict:
            snake_case_   			=   posixpath.join(self.local_root_dir,				self._strip_protocol(lowerCAmelCase__))
            snake_case_   			=   dict(self._fs.info(lowerCAmelCase__,				*lowerCAmelCase__,				**lowerCAmelCase__))
            snake_case_   			=   out['name'][len(self.local_root_dir) :]
            return out







       def 						a_		(					self,				lowerCAmelCase__,				lowerCAmelCase__,				*lowerCAmelCase__,				**lowerCAmelCase__) -> str:
            snake_case_   			=   posixpath.join(self.local_root_dir,				self._strip_protocol(lowerCAmelCase__))
            snake_case_   			=   posixpath.join(self.local_root_dir,				self._strip_protocol(lowerCAmelCase__))
            return self._fs.cp_file(lowerCAmelCase__,				lowerCAmelCase__,				*lowerCAmelCase__,				**lowerCAmelCase__)







       def 						a_		(					self,				lowerCAmelCase__,				*lowerCAmelCase__,				**lowerCAmelCase__) -> Union[str, Any]:
            snake_case_   			=   posixpath.join(self.local_root_dir,				self._strip_protocol(lowerCAmelCase__))
            return self._fs.rm_file(lowerCAmelCase__,				*lowerCAmelCase__,				**lowerCAmelCase__)







       def 						a_		(					self,				lowerCAmelCase__,				*lowerCAmelCase__,				**lowerCAmelCase__) -> List[Any]:
            snake_case_   			=   posixpath.join(self.local_root_dir,				self._strip_protocol(lowerCAmelCase__))
            return self._fs.rm(lowerCAmelCase__,				*lowerCAmelCase__,				**lowerCAmelCase__)







       def 						a_		(					self,				lowerCAmelCase__,				*lowerCAmelCase__,				**lowerCAmelCase__) -> Optional[Any]:
            snake_case_   			=   posixpath.join(self.local_root_dir,				self._strip_protocol(lowerCAmelCase__))
            return self._fs._open(lowerCAmelCase__,				*lowerCAmelCase__,				**lowerCAmelCase__)







       def 						a_		(					self,				lowerCAmelCase__) -> Tuple:
            snake_case_   			=   posixpath.join(self.local_root_dir,				self._strip_protocol(lowerCAmelCase__))
            return self._fs.created(lowerCAmelCase__)







       def 						a_		(					self,				lowerCAmelCase__) -> int:
            snake_case_   			=   posixpath.join(self.local_root_dir,				self._strip_protocol(lowerCAmelCase__))
            return self._fs.modified(lowerCAmelCase__)







       @classmethod
       def 						a_		(					cls,				lowerCAmelCase__) -> Any:
            snake_case_   			=   stringify_path(lowerCAmelCase__)
            if path.startswith('mock://'):
                 snake_case_   			=   path[7:]
            return path




class 					UpperCamelCase   (							lowerCAmelCase__  ):
       SCREAMING_SNAKE_CASE_						      = "tmp"
       SCREAMING_SNAKE_CASE_						      = None







       def __init__(					self,				*lowerCAmelCase__,				**lowerCAmelCase__) -> str:
            assert self.tmp_dir is not None, "TmpDirFileSystem.tmp_dir is not set"
            super().__init__(*lowerCAmelCase__,				**lowerCAmelCase__,				local_root_dir=self.tmp_dir,				auto_mkdir=lowerCAmelCase__)







       @classmethod
       def 						a_		(					cls,				lowerCAmelCase__) -> List[str]:
            snake_case_   			=   stringify_path(lowerCAmelCase__)
            if path.startswith('tmp://'):
                 snake_case_   			=   path[6:]
            return path







@pytest.fixture
def    UpperCAmelCase       (						)							-> Tuple:
     snake_case_   			=   MockFileSystem
     snake_case_   			=   TmpDirFileSystem
     yield
     del _fsspec_registry["mock"]
     del _fsspec_registry["tmp"]







@pytest.fixture
def    UpperCAmelCase       (						UpperCAmelCase		,   UpperCAmelCase   )							-> Tuple:
     snake_case_   			=   tmp_path_factory.mktemp('mockfs'   )
     return MockFileSystem(local_root_dir=UpperCAmelCase		,   auto_mkdir=UpperCAmelCase   )







@pytest.fixture
def    UpperCAmelCase       (						UpperCAmelCase		,   UpperCAmelCase   )							-> int:
     snake_case_   			=   tmp_path_factory.mktemp('tmpfs'   )
     with patch.object(UpperCAmelCase		,   'tmp_dir'		,   UpperCAmelCase   ):
          yield TmpDirFileSystem()

