

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







 """simple docstring"""







 def __init__(     self				:     Union[str, Any]      ):
    self.register_to_config()
    super().__init__()


 def    _lowercase					(     self				:     List[Any],   UpperCAmelCase__				:     int,   UpperCAmelCase__				:     List[str]      ):




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


 def    _lowercase					(     self				:     int,   UpperCAmelCase__				:     Union[str, Any]      ):
    __lowercase								=		[]
    __lowercase								=		config_dict.copy()
    for key in config_dict.keys():
       if key.startswith("_"      ):
          temp_dict.pop(__SCREAMING_SNAKE_CASE      )
          meta_keys.append(__SCREAMING_SNAKE_CASE      )
    return (temp_dict, meta_keys)


 @torch.no_grad()
 def    _lowercase					(     self				:     Tuple,   UpperCAmelCase__				:     Any,   **UpperCAmelCase__				:     Union[str, Any]      ):
    __lowercase								=		kwargs.pop("cache_dir",   __SCREAMING_SNAKE_CASE      )
    __lowercase								=		kwargs.pop("resume_download",   __SCREAMING_SNAKE_CASE      )
    __lowercase								=		kwargs.pop("force_download",   __SCREAMING_SNAKE_CASE      )
    __lowercase								=		kwargs.pop("proxies",   __SCREAMING_SNAKE_CASE      )
    __lowercase								=		kwargs.pop("local_files_only",   __SCREAMING_SNAKE_CASE      )
    __lowercase								=		kwargs.pop("use_auth_token",   __SCREAMING_SNAKE_CASE      )
    __lowercase								=		kwargs.pop("revision",   __SCREAMING_SNAKE_CASE      )
    __lowercase								=		kwargs.pop("torch_dtype",   __SCREAMING_SNAKE_CASE      )
    __lowercase								=		kwargs.pop("device_map",   __SCREAMING_SNAKE_CASE      )

    __lowercase								=		kwargs.pop("alpha",   0.5      )
    __lowercase								=		kwargs.pop("interp",   __SCREAMING_SNAKE_CASE      )

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

    __lowercase								=		len(__SCREAMING_SNAKE_CASE      )
    # Ignore result from model_index_json comparision of the two checkpoints
    __lowercase								=		kwargs.pop("force",   __SCREAMING_SNAKE_CASE      )

    # 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 '_'
    __lowercase								=		[]
    for pretrained_model_name_or_path in pretrained_model_name_or_path_list:
       __lowercase								=		DiffusionPipeline.load_config(
           __SCREAMING_SNAKE_CASE,   cache_dir=__SCREAMING_SNAKE_CASE,   resume_download=__SCREAMING_SNAKE_CASE,   force_download=__SCREAMING_SNAKE_CASE,   proxies=__SCREAMING_SNAKE_CASE,   local_files_only=__SCREAMING_SNAKE_CASE,   use_auth_token=__SCREAMING_SNAKE_CASE,   revision=__SCREAMING_SNAKE_CASE,   )
       config_dicts.append(__SCREAMING_SNAKE_CASE      )

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

       __lowercase								=		(
           pretrained_model_name_or_path
           if os.path.isdir(__SCREAMING_SNAKE_CASE      )
           else snapshot_download(
               __SCREAMING_SNAKE_CASE,   cache_dir=__SCREAMING_SNAKE_CASE,   resume_download=__SCREAMING_SNAKE_CASE,   proxies=__SCREAMING_SNAKE_CASE,   local_files_only=__SCREAMING_SNAKE_CASE,   use_auth_token=__SCREAMING_SNAKE_CASE,   revision=__SCREAMING_SNAKE_CASE,   allow_patterns=__SCREAMING_SNAKE_CASE,   user_agent=__SCREAMING_SNAKE_CASE,   )
       )
       print("Cached Folder",   __SCREAMING_SNAKE_CASE      )
       cached_folders.append(__SCREAMING_SNAKE_CASE      )

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

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

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

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

             __lowercase								=		getattr(__SCREAMING_SNAKE_CASE,   "load_state_dict"      )
             __lowercase								=		(
                 safetensors.torch.load_file(__SCREAMING_SNAKE_CASE      )
                 if (is_safetensors_available() and checkpoint_path_a.endswith(".safetensors"      ))
                 else torch.load(__SCREAMING_SNAKE_CASE,   map_location="cpu"      )
             )
             __lowercase								=		None
             if checkpoint_path_a:
                __lowercase								=		(
                    safetensors.torch.load_file(__SCREAMING_SNAKE_CASE      )
                    if (is_safetensors_available() and checkpoint_path_a.endswith(".safetensors"      ))
                    else torch.load(__SCREAMING_SNAKE_CASE,   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(__SCREAMING_SNAKE_CASE      )}"""      )
             continue
          print(F"""MERGING {attr}"""      )

          for key in theta_a.keys():
             if theta_a:
                __lowercase								=		theta_func(theta_a[key],   theta_a[key],   theta_a[key],   __SCREAMING_SNAKE_CASE      )
             else:
                __lowercase								=		theta_func(theta_a[key],   theta_a[key],   __SCREAMING_SNAKE_CASE,   __SCREAMING_SNAKE_CASE      )

          del theta_a
          del theta_a
          update_theta_a(__SCREAMING_SNAKE_CASE      )

          del theta_a
    return final_pipe


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


 @staticmethod
 def    _lowercase					(     UpperCAmelCase__				:     Any,   UpperCAmelCase__				:     Dict,   UpperCAmelCase__				:     Union[str, Any],   UpperCAmelCase__				:     Union[str, Any]      ):
    __lowercase								=		alpha * alpha * (3 - (2 * alpha))
    return thetaa + ((thetaa - thetaa) * alpha)


 @staticmethod
 def    _lowercase					(     UpperCAmelCase__				:     Any,   UpperCAmelCase__				:     Tuple,   UpperCAmelCase__				:     Any,   UpperCAmelCase__				:     List[Any]      ):
    import math

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



 @staticmethod
 def    _lowercase					(     UpperCAmelCase__				:     List[str],   UpperCAmelCase__				:     Dict,   UpperCAmelCase__				:     int,   UpperCAmelCase__				:     Any      ):
    return thetaa + (thetaa - thetaa) * (1.0 - alpha)








import tempfile
import unittest

import torch

from diffusers import UNetaDConditionModel
from diffusers.training_utils import EMAModel
from diffusers.utils.testing_utils import enable_full_determinism, skip_mps, torch_device


enable_full_determinism()



class  UpperCAmelCase_						(				unittest.TestCase):
       '''simple docstring'''
       __UpperCamelCase   :  Tuple				=       "hf-internal-testing/tiny-stable-diffusion-pipe"
       __UpperCamelCase   :  Optional[Any]				=       1
       __UpperCamelCase   :  str				=       77
       __UpperCamelCase   :  Tuple				=       32
       __UpperCamelCase   :  List[str]				=       4
       __UpperCamelCase   :  Optional[int]				=       64
       __UpperCamelCase   :  Union[str, Any]				=       torch.manual_seed(0)






       def 			_lowercase (  self			,       __SCREAMING_SNAKE_CASE=0.9_999    ):




              """simple docstring"""




              UpperCamelCase     :						Optional[int] 						=		UNetaDConditionModel.from_pretrained(self.model_id			,       subfolder='''unet'''    )
              UpperCamelCase     :						Any 						=		unet.to(__SCREAMING_SNAKE_CASE    )
              UpperCamelCase     :						Dict 						=		EMAModel(unet.parameters()			,       decay=__SCREAMING_SNAKE_CASE			,       model_cls=__SCREAMING_SNAKE_CASE			,       model_config=unet.config    )
              return unet, ema_unet






       def 			_lowercase (  self    ):




              """simple docstring"""




              UpperCamelCase     :						Tuple 						=		torch.randn(
                  self.batch_size			,       self.num_in_channels			,       self.latent_height			,       self.latent_width			,       generator=self.generator    ).to(__SCREAMING_SNAKE_CASE    )
              UpperCamelCase     :						Optional[int] 						=		torch.randint(0			,       1_000			,       size=(self.batch_size,)			,       generator=self.generator    ).to(__SCREAMING_SNAKE_CASE    )
              UpperCamelCase     :						List[Any] 						=		torch.randn(
                  self.batch_size			,       self.prompt_length			,       self.text_encoder_hidden_dim			,       generator=self.generator    ).to(__SCREAMING_SNAKE_CASE    )
              return noisy_latents, timesteps, encoder_hidden_states






       def 			_lowercase (  self			,       __SCREAMING_SNAKE_CASE    ):




              """simple docstring"""




              UpperCamelCase     :						Dict 						=		{}
              for k, param in unet.state_dict().items():
                     UpperCamelCase     :						int 						=		torch.randn_like(__SCREAMING_SNAKE_CASE    ) + (param * torch.randn_like(__SCREAMING_SNAKE_CASE    ))
                     updated_state_dict.update({k: updated_param}    )
              unet.load_state_dict(__SCREAMING_SNAKE_CASE    )
              return unet






       def 			_lowercase (  self    ):




              """simple docstring"""




              UpperCamelCase						,			UpperCamelCase     :						Optional[Any] 						=		self.get_models()
              # Take the first (hypothetical) EMA step.
              ema_unet.step(unet.parameters()    )
              assert ema_unet.optimization_step == 1

              # Take two more.
              for _ in range(2    ):
                     ema_unet.step(unet.parameters()    )
              assert ema_unet.optimization_step == 3






       def 			_lowercase (  self    ):




              """simple docstring"""




              UpperCamelCase						,			UpperCamelCase     :						List[Any] 						=		self.get_models()
              # Since the `unet` is not being updated (i.e., backprop'd)
              # there won't be any difference between the `params` of `unet`
              # and `ema_unet` even if we call `ema_unet.step(unet.parameters())`.
              ema_unet.step(unet.parameters()    )
              UpperCamelCase     :						int 						=		list(unet.parameters()    )
              for s_param, param in zip(ema_unet.shadow_params			,       __SCREAMING_SNAKE_CASE    ):
                     assert torch.allclose(__SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE    )

              # The above holds true even if we call `ema.step()` multiple times since
              # `unet` params are still not being updated.
              for _ in range(4    ):
                     ema_unet.step(unet.parameters()    )
              for s_param, param in zip(ema_unet.shadow_params			,       __SCREAMING_SNAKE_CASE    ):
                     assert torch.allclose(__SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE    )






       def 			_lowercase (  self    ):




              """simple docstring"""




              UpperCamelCase						,			UpperCamelCase     :						Any 						=		self.get_models()
              # Here we simulate the parameter updates for `unet`. Since there might
              # be some parameters which are initialized to zero we take extra care to
              # initialize their values to something non-zero before the multiplication.
              UpperCamelCase     :						Union[str, Any] 						=		self.simulate_backprop(__SCREAMING_SNAKE_CASE    )

              # Take the EMA step.
              ema_unet.step(unet_pseudo_updated_step_one.parameters()    )

              # Now the EMA'd parameters won't be equal to the original model parameters.
              UpperCamelCase     :						Any 						=		list(unet_pseudo_updated_step_one.parameters()    )
              for s_param, param in zip(ema_unet.shadow_params			,       __SCREAMING_SNAKE_CASE    ):
                     assert ~torch.allclose(__SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE    )

              # Ensure this is the case when we take multiple EMA steps.
              for _ in range(4    ):
                     ema_unet.step(unet.parameters()    )
              for s_param, param in zip(ema_unet.shadow_params			,       __SCREAMING_SNAKE_CASE    ):
                     assert ~torch.allclose(__SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE    )






       def 			_lowercase (  self    ):




              """simple docstring"""




              UpperCamelCase						,			UpperCamelCase     :						Optional[int] 						=		self.get_models()

              # First backprop + EMA
              UpperCamelCase     :						Dict 						=		self.simulate_backprop(__SCREAMING_SNAKE_CASE    )
              ema_unet.step(unet_step_one.parameters()    )
              UpperCamelCase     :						List[str] 						=		ema_unet.shadow_params

              # Second backprop + EMA
              UpperCamelCase     :						str 						=		self.simulate_backprop(__SCREAMING_SNAKE_CASE    )
              ema_unet.step(unet_step_two.parameters()    )
              UpperCamelCase     :						List[Any] 						=		ema_unet.shadow_params

              for step_one, step_two in zip(__SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE    ):
                     assert ~torch.allclose(__SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE    )






       def 			_lowercase (  self    ):




              """simple docstring"""




              UpperCamelCase						,			UpperCamelCase     :						int 						=		self.get_models(decay=0.0    )
              UpperCamelCase     :						Optional[int] 						=		self.simulate_backprop(__SCREAMING_SNAKE_CASE    )
              ema_unet.step(unet_step_one.parameters()    )
              UpperCamelCase     :						Union[str, Any] 						=		ema_unet.shadow_params

              UpperCamelCase     :						Any 						=		self.simulate_backprop(__SCREAMING_SNAKE_CASE    )
              ema_unet.step(unet_step_two.parameters()    )
              UpperCamelCase     :						int 						=		ema_unet.shadow_params

              for step_one, step_two in zip(__SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE    ):
                     assert torch.allclose(__SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE    )






       @skip_mps
       def 			_lowercase (  self    ):




              """simple docstring"""




              UpperCamelCase						,			UpperCamelCase     :						str 						=		self.get_models()
              UpperCamelCase						,			UpperCamelCase						,			UpperCamelCase     :						Union[str, Any] 						=		self.get_dummy_inputs()

              with tempfile.TemporaryDirectory() as tmpdir:
                     ema_unet.save_pretrained(__SCREAMING_SNAKE_CASE    )
                     UpperCamelCase     :						Optional[int] 						=		UNetaDConditionModel.from_pretrained(__SCREAMING_SNAKE_CASE			,       model_cls=__SCREAMING_SNAKE_CASE    )
                     UpperCamelCase     :						List[str] 						=		loaded_unet.to(unet.device    )

              # Since no EMA step has been performed the outputs should match.
              UpperCamelCase     :						Optional[Any] 						=		unet(__SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE    ).sample
              UpperCamelCase     :						List[str] 						=		loaded_unet(__SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE    ).sample

              assert torch.allclose(__SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE			,       atol=1e-4    )

