

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      snake_case_       (		__A ):







			'''simple docstring'''

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







			def     snake_case__(					self			:    List[str]							,       _UpperCamelCase			:    Optional[Any]							,       _UpperCamelCase			:    Union[str, Any]							)				->Optional[int]:
									if dicta == dicta:
															return True
									else:
															snake_case_,   snake_case_       =     self._remove_meta_keys(_UpperCamelCase							)
															snake_case_,   snake_case_       =     self._remove_meta_keys(_UpperCamelCase							)
															if configa == configa:
																					print(f'''Warning !: Mismatch in keys {meta_keysa} and {meta_keysa}.'''							)
																					return True
									return False







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







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

									snake_case_       =     kwargs.pop('''alpha'''							,       0.5							)
									snake_case_       =     kwargs.pop('''interp'''							,       _UpperCamelCase							)

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

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

									# 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(
															    _UpperCamelCase							,       cache_dir=_UpperCamelCase							,       resume_download=_UpperCamelCase							,       force_download=_UpperCamelCase							,       proxies=_UpperCamelCase							,       local_files_only=_UpperCamelCase							,       use_auth_token=_UpperCamelCase							,       revision=_UpperCamelCase							,       )
															config_dicts.append(_UpperCamelCase							)

									snake_case_       =     True
									for idx in range(1							,       len(_UpperCamelCase							)							):
															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(_UpperCamelCase							,       _UpperCamelCase							):
															snake_case_       =     [k for k in config_dict.keys() if not k.startswith('''_'''							)]
															snake_case_       =     [os.path.join(_UpperCamelCase							,       '''*'''							) 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(_UpperCamelCase							)
															    else snapshot_download(
															        _UpperCamelCase							,       cache_dir=_UpperCamelCase							,       resume_download=_UpperCamelCase							,       proxies=_UpperCamelCase							,       local_files_only=_UpperCamelCase							,       use_auth_token=_UpperCamelCase							,       revision=_UpperCamelCase							,       allow_patterns=_UpperCamelCase							,       user_agent=_UpperCamelCase							,       )
															)
															print('''Cached Folder'''							,       _UpperCamelCase							)
															cached_folders.append(_UpperCamelCase							)

									# 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=_UpperCamelCase							,       device_map=_UpperCamelCase							)
									final_pipe.to(self.device							)

									snake_case_       =     None
									if len(_UpperCamelCase							) > 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]							,       _UpperCamelCase							)
																					if os.path.exists(_UpperCamelCase							):
																											snake_case_       =     [
																											    *glob.glob(os.path.join(_UpperCamelCase							,       '''*.safetensors'''							)							),
																											    *glob.glob(os.path.join(_UpperCamelCase							,       '''*.bin'''							)							),
																											]
																											snake_case_       =     files[0] if len(_UpperCamelCase							) > 0 else None
																					if len(_UpperCamelCase							) < 3:
																											snake_case_       =     None
																					else:
																											snake_case_       =     os.path.join(cached_folders[2]							,       _UpperCamelCase							)
																											if os.path.exists(_UpperCamelCase							):
																																	snake_case_       =     [
																																	    *glob.glob(os.path.join(_UpperCamelCase							,       '''*.safetensors'''							)							),
																																	    *glob.glob(os.path.join(_UpperCamelCase							,       '''*.bin'''							)							),
																																	]
																																	snake_case_       =     files[0] if len(_UpperCamelCase							) > 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(_UpperCamelCase							,       _UpperCamelCase							)
																											if isinstance(_UpperCamelCase							,       _UpperCamelCase							):  # ignore requires_safety_checker boolean
																																	continue
																											snake_case_       =     getattr(_UpperCamelCase							,       '''state_dict'''							)
																											snake_case_       =     theta_a()

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

																					del theta_a
																					del theta_a
																					update_theta_a(_UpperCamelCase							)

																					del theta_a
									return final_pipe







			@staticmethod
			def     snake_case__(					_UpperCamelCase			:    List[str]							,       _UpperCamelCase			:    Tuple							,       _UpperCamelCase			:    List[Any]							,       _UpperCamelCase			:    str							)				->List[str]:
									return ((1 - alpha) * thetaa) + (alpha * thetaa)







			@staticmethod
			def     snake_case__(					_UpperCamelCase			:    Union[str, Any]							,       _UpperCamelCase			:    List[str]							,       _UpperCamelCase			:    Optional[int]							,       _UpperCamelCase			:    str							)				->int:
									snake_case_       =     alpha * alpha * (3 - (2 * alpha))
									return thetaa + ((thetaa - thetaa) * alpha)







			@staticmethod
			def     snake_case__(					_UpperCamelCase			:    List[str]							,       _UpperCamelCase			:    List[str]							,       _UpperCamelCase			:    Union[str, Any]							,       _UpperCamelCase			:    Any							)				->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     snake_case__(					_UpperCamelCase			:    List[Any]							,       _UpperCamelCase			:    Optional[int]							,       _UpperCamelCase			:    str							,       _UpperCamelCase			:    Union[str, Any]							)				->List[Any]:
									return thetaa + (thetaa - thetaa) * (1.0 - alpha)


import json
import os
from collections.abc import Mapping
from typing import Dict, List, Optional, Union

import numpy as np
import torch

from .imports import is_safetensors_available


def   __SCREAMING_SNAKE_CASE (SCREAMING_SNAKE_CASE__  ,		SCREAMING_SNAKE_CASE__  ,		SCREAMING_SNAKE_CASE__  ,		SCREAMING_SNAKE_CASE__=None	):
						snake_case_       =     None
						# Check the string instead of the dtype to be compatible with versions of PyTorch that don't have bfloat16.
						if str(weight.dtype	) == "torch.bfloat16":
												# Need to reinterpret the underlined data as int16 since NumPy does not handle bfloat16s.
												snake_case_       =     weight.view(torch.intaa	)
												snake_case_       =     '''bfloat16'''
						snake_case_       =     weight.cpu().numpy()
						snake_case_       =     os.path.join(SCREAMING_SNAKE_CASE__  ,		F'''{weight_name}.dat'''	)
						if index is not None:
												if dtype is None:
																		snake_case_       =     str(array.dtype	)
												snake_case_       =     {'''dtype''': dtype, '''shape''': list(array.shape	)}
						if array.ndim == 0:
												snake_case_       =     array[None]
						snake_case_       =     np.memmap(SCREAMING_SNAKE_CASE__  ,		dtype=array.dtype  ,		mode='''w+'''  ,		shape=array.shape	)
						snake_case_       =     array[:]
						file_array.flush()
						return index


def   __SCREAMING_SNAKE_CASE (SCREAMING_SNAKE_CASE__  ,		SCREAMING_SNAKE_CASE__	):
						snake_case_       =     tuple(weight_info['''shape''']	)
						if shape == ():
												# NumPy memory-mapped arrays can't have 0 dims so it was saved as 1d tensor
												snake_case_       =     (1,)

						snake_case_       =     weight_info['''dtype''']
						if dtype == "bfloat16":
												# NumPy does not support bfloat16 so this was saved as a int16
												snake_case_       =     '''int16'''

						snake_case_       =     np.memmap(SCREAMING_SNAKE_CASE__  ,		dtype=SCREAMING_SNAKE_CASE__  ,		shape=SCREAMING_SNAKE_CASE__  ,		mode='''r'''	)

						if len(weight_info['''shape''']	) == 0:
												snake_case_       =     weight[0]
						snake_case_       =     torch.tensor(SCREAMING_SNAKE_CASE__	)
						if weight_info["dtype"] == "bfloat16":
												snake_case_       =     weight.view(torch.bfloataa	)

						return weight


def   __SCREAMING_SNAKE_CASE (SCREAMING_SNAKE_CASE__  ,		SCREAMING_SNAKE_CASE__	):
						if index is None or len(SCREAMING_SNAKE_CASE__	) == 0:
												# Nothing to save
												return

						snake_case_       =     os.path.join(SCREAMING_SNAKE_CASE__  ,		'''index.json'''	)
						if os.path.isfile(SCREAMING_SNAKE_CASE__	):
												with open(SCREAMING_SNAKE_CASE__  ,		'''r'''  ,		encoding='''utf-8'''	) as f:
																		snake_case_       =     json.load(SCREAMING_SNAKE_CASE__	)
						else:
												snake_case_       =     {}
						current_index.update(SCREAMING_SNAKE_CASE__	)

						with open(SCREAMING_SNAKE_CASE__  ,		'''w'''  ,		encoding='''utf-8'''	) as f:
												json.dump(SCREAMING_SNAKE_CASE__  ,		SCREAMING_SNAKE_CASE__  ,		indent=2	)


def   __SCREAMING_SNAKE_CASE (SCREAMING_SNAKE_CASE__  ,		SCREAMING_SNAKE_CASE__	):
						os.makedirs(SCREAMING_SNAKE_CASE__  ,		exist_ok=SCREAMING_SNAKE_CASE__	)
						snake_case_       =     {}
						for name, parameter in state_dict.items():
												snake_case_       =     offload_weight(SCREAMING_SNAKE_CASE__  ,		SCREAMING_SNAKE_CASE__  ,		SCREAMING_SNAKE_CASE__  ,		index=SCREAMING_SNAKE_CASE__	)

						# Update index
						save_offload_index(SCREAMING_SNAKE_CASE__  ,		SCREAMING_SNAKE_CASE__	)





class      snake_case_       (		__A ):







			'''simple docstring'''

			def __init__(					self			:    Optional[Any]							,       _UpperCamelCase			:    Mapping							,       _UpperCamelCase			:    str							)				->int:
									snake_case_       =     dataset
									snake_case_       =     prefix







			def __getitem__(					self			:    List[Any]							,       _UpperCamelCase			:    Any							)				->Tuple:
									return self.dataset[f'''{self.prefix}{key}''']







			def __iter__(					self			:    Optional[int]							)				->int:
									return iter([key for key in self.dataset if key.startswith(self.prefix							)]							)







			def __len__(					self			:    Dict							)				->int:
									return len(self.dataset							)






class      snake_case_       (		__A ):







			'''simple docstring'''

			def __init__(					self			:    str							,       _UpperCamelCase			:    Dict[str, torch.Tensor] = None							,       _UpperCamelCase			:    Optional[Union[str, os.PathLike]] = None							,       _UpperCamelCase			:    Mapping = None							,       _UpperCamelCase			:    Tuple=None							,       )				->Dict:
									if state_dict is None and save_folder is None:
															raise ValueError('''Need either a `state_dict` or a `save_folder` containing offloaded weights.'''							)

									snake_case_       =     {} if state_dict is None else state_dict
									snake_case_       =     save_folder
									if index is None and save_folder is not None:
															with open(os.path.join(_UpperCamelCase							,       '''index.json'''							)							) as f:
																					snake_case_       =     json.load(_UpperCamelCase							)
									snake_case_       =     {} if index is None else index
									snake_case_       =     list(self.state_dict.keys()							)
									self.all_keys.extend([key for key in self.index if key not in self.all_keys]							)
									snake_case_       =     device







			def __getitem__(					self			:    int							,       _UpperCamelCase			:    str							)				->Union[str, Any]:
									# State dict gets priority
									if key in self.state_dict:
															return self.state_dict[key]
									snake_case_       =     self.index[key]
									if weight_info.get('''safetensors_file'''							) is not None:
															if not is_safetensors_available():
																					raise ImportError('''These offloaded weights require the use of safetensors: `pip install safetensors`.'''							)

															from safetensors import safe_open

															snake_case_       =     '''cpu''' if self.device is None else self.device
															with safe_open(weight_info['''safetensors_file''']							,       framework='''pt'''							,       device=_UpperCamelCase							) as f:
																					snake_case_       =     f.get_tensor(weight_info.get('''weight_name'''							,       _UpperCamelCase							)							)

															if "dtype" in weight_info:
																					return tensor.to(getattr(_UpperCamelCase							,       weight_info['''dtype''']							)							)
															else:
																					return tensor

									snake_case_       =     os.path.join(self.save_folder							,       f'''{key}.dat'''							)
									return load_offloaded_weight(_UpperCamelCase							,       _UpperCamelCase							)







			def __iter__(					self			:    int							)				->List[Any]:
									return iter(self.all_keys							)







			def __len__(					self			:    Optional[Any]							)				->Optional[Any]:
									return len(self.all_keys							)






def   __SCREAMING_SNAKE_CASE (SCREAMING_SNAKE_CASE__  ,		SCREAMING_SNAKE_CASE__	):
						snake_case_       =     {}
						for module_name in submodule_names:
												# We want to catch module_name parameter (module_name.xxx) or potentially module_name, but not any of the
												# submodules that could being like module_name (transformers.h.1 and transformers.h.10 for instance)
												result.update(
												    {
												        key: param
												        for key, param in state_dict.items()
												        if key == module_name or key.startswith(module_name + '''.'''	)
												    }	)
						return result