






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


   '''simple docstring'''

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

     '''simple docstring'''
     self.register_to_config()
     super().__init__()


   def 							_SCREAMING_SNAKE_CASE							(      self					:	Optional[int],			_lowerCamelCase					:	Tuple,			_lowerCamelCase					:	List[str]    ):

     '''simple docstring'''
     if dicta == dicta:
       return True
     else:
       __A     ,       __A		     =					self._remove_meta_keys(_lowerCamelCase    )
       __A     ,       __A		     =					self._remove_meta_keys(_lowerCamelCase    )
       if configa == configa:
         print(f'Warning !: Mismatch in keys {meta_keysa} and {meta_keysa}.'    )
         return True
     return False


   def 							_SCREAMING_SNAKE_CASE							(      self					:	Any,			_lowerCamelCase					:	Dict    ):

     '''simple docstring'''
     __A		     =					[]
     __A		     =					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 							_SCREAMING_SNAKE_CASE							(      self					:	int,			_lowerCamelCase					:	List[Union[str, os.PathLike]],			**_lowerCamelCase					:	Optional[int]    ):

     '''simple docstring'''
     __A		     =					kwargs.pop('''cache_dir''',			_lowerCamelCase    )
     __A		     =					kwargs.pop('''resume_download''',			_lowerCamelCase    )
     __A		     =					kwargs.pop('''force_download''',			_lowerCamelCase    )
     __A		     =					kwargs.pop('''proxies''',			_lowerCamelCase    )
     __A		     =					kwargs.pop('''local_files_only''',			_lowerCamelCase    )
     __A		     =					kwargs.pop('''use_auth_token''',			_lowerCamelCase    )
     __A		     =					kwargs.pop('''revision''',			_lowerCamelCase    )
     __A		     =					kwargs.pop('''torch_dtype''',			_lowerCamelCase    )
     __A		     =					kwargs.pop('''device_map''',			_lowerCamelCase    )

     __A		     =					kwargs.pop('''alpha''',			0.5    )
     __A		     =					kwargs.pop('''interp''',			_lowerCamelCase    )

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

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

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

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

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

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

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

           __A		     =					getattr(_lowerCamelCase,			'''load_state_dict'''    )
           __A		     =					(
               safetensors.torch.load_file(_lowerCamelCase    )
               if (is_safetensors_available() and checkpoint_path_a.endswith('''.safetensors'''    ))
               else torch.load(_lowerCamelCase,			map_location='''cpu'''    )
           )
           __A		     =					None
           if checkpoint_path_a:
             __A		     =					(
                 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:
             __A		     =					theta_func(theta_a[key],			theta_a[key],			theta_a[key],			_lowerCamelCase    )
           else:
             __A		     =					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 							_SCREAMING_SNAKE_CASE							(      _lowerCamelCase					:	Optional[Any],			_lowerCamelCase					:	Optional[Any],			_lowerCamelCase					:	List[str],			_lowerCamelCase					:	Any    ):

     '''simple docstring'''
     return ((1 - alpha) * thetaa) + (alpha * thetaa)


   @staticmethod
   def 							_SCREAMING_SNAKE_CASE							(      _lowerCamelCase					:	Optional[Any],			_lowerCamelCase					:	List[Any],			_lowerCamelCase					:	Optional[int],			_lowerCamelCase					:	Union[str, Any]    ):

     '''simple docstring'''
     __A		     =					alpha * alpha * (3 - (2 * alpha))
     return thetaa + ((thetaa - thetaa) * alpha)


   @staticmethod
   def 							_SCREAMING_SNAKE_CASE							(      _lowerCamelCase					:	Tuple,			_lowerCamelCase					:	Tuple,			_lowerCamelCase					:	Any,			_lowerCamelCase					:	Dict    ):

     '''simple docstring'''
     import math

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






   @staticmethod
   def 							_SCREAMING_SNAKE_CASE							(      _lowerCamelCase					:	List[str],			_lowerCamelCase					:	Tuple,			_lowerCamelCase					:	Any,			_lowerCamelCase					:	int    ):

     '''simple docstring'''
     return thetaa + (thetaa - thetaa) * (1.0 - alpha)












"""simple docstring"""




import inspect
import unittest

from transformers import SwinaSRConfig
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import is_torch_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
     import torch
     from torch import nn

     from transformers import SwinaSRForImageSuperResolution, SwinaSRModel
     from transformers.models.swinasr.modeling_swinasr import SWIN2SR_PRETRAINED_MODEL_ARCHIVE_LIST

if is_vision_available():
     from PIL import Image

     from transformers import SwinaSRImageProcessor




class  lowercase_     :






      def __init__(     self  ,  a_  ,  a_=1_3  ,  a_=3_2  ,  a_=1  ,  a_=3  ,  a_=1_6  ,  a_=[1, 2, 1]  ,  a_=[2, 2, 4]  ,  a_=2  ,  a_=2.0  ,  a_=True  ,  a_=0.0  ,  a_=0.0  ,  a_=0.1  ,  a_="gelu"  ,  a_=False  ,  a_=True  ,  a_=0.02  ,  a_=1E-5  ,  a_=True  ,  a_=None  ,  a_=False  ,  a_=2  ,  )						->str:







        '''simple docstring'''


        _a					      =						parent
        _a					      =						batch_size
        _a					      =						image_size
        _a					      =						patch_size
        _a					      =						num_channels
        _a					      =						embed_dim
        _a					      =						depths
        _a					      =						num_heads
        _a					      =						window_size
        _a					      =						mlp_ratio
        _a					      =						qkv_bias
        _a					      =						hidden_dropout_prob
        _a					      =						attention_probs_dropout_prob
        _a					      =						drop_path_rate
        _a					      =						hidden_act
        _a					      =						use_absolute_embeddings
        _a					      =						patch_norm
        _a					      =						layer_norm_eps
        _a					      =						initializer_range
        _a					      =						is_training
        _a					      =						scope
        _a					      =						use_labels
        _a					      =						upscale

        # here we set some attributes to make tests pass
        _a					      =						len(a_						)
        _a					      =						embed_dim
        _a					      =						(image_size // patch_size) ** 2






      def       lowerCamelCase__     (     self						)						->Any:







        '''simple docstring'''


        _a					      =						floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]						)

        _a					      =						None
        if self.use_labels:
          _a					      =						ids_tensor([self.batch_size]  ,  self.type_sequence_label_size						)

        _a					      =						self.get_config()

        return config, pixel_values, labels






      def       lowerCamelCase__     (     self						)						->Union[str, Any]:







        '''simple docstring'''


        return SwinaSRConfig(
            image_size=self.image_size  ,  patch_size=self.patch_size  ,  num_channels=self.num_channels  ,  embed_dim=self.embed_dim  ,  depths=self.depths  ,  num_heads=self.num_heads  ,  window_size=self.window_size  ,  mlp_ratio=self.mlp_ratio  ,  qkv_bias=self.qkv_bias  ,  hidden_dropout_prob=self.hidden_dropout_prob  ,  attention_probs_dropout_prob=self.attention_probs_dropout_prob  ,  drop_path_rate=self.drop_path_rate  ,  hidden_act=self.hidden_act  ,  use_absolute_embeddings=self.use_absolute_embeddings  ,  path_norm=self.patch_norm  ,  layer_norm_eps=self.layer_norm_eps  ,  initializer_range=self.initializer_range  ,  upscale=self.upscale  ,  )






      def       lowerCamelCase__     (     self  ,  a_  ,  a_  ,  a_						)						->Union[str, Any]:







        '''simple docstring'''


        _a					      =						SwinaSRModel(config=a_						)
        model.to(a_						)
        model.eval()
        _a					      =						model(a_						)

        self.parent.assertEqual(
            result.last_hidden_state.shape  ,  (self.batch_size, self.embed_dim, self.image_size, self.image_size)						)






      def       lowerCamelCase__     (     self  ,  a_  ,  a_  ,  a_						)						->Optional[int]:







        '''simple docstring'''


        _a					      =						SwinaSRForImageSuperResolution(a_						)
        model.to(a_						)
        model.eval()
        _a					      =						model(a_						)

        _a					      =						self.image_size * self.upscale
        self.parent.assertEqual(
            result.reconstruction.shape  ,  (self.batch_size, self.num_channels, expected_image_size, expected_image_size)						)






      def       lowerCamelCase__     (     self						)						->Tuple:







        '''simple docstring'''


        _a					      =						self.prepare_config_and_inputs()
        _a      ,						_a      ,						_a					      =						config_and_inputs
        _a					      =						{"pixel_values": pixel_values}
        return config, inputs_dict







@require_torch
class  lowercase_     (_UpperCAmelCase,     _UpperCAmelCase,     unittest.TestCase  ):
      A__						:		str							      =					(SwinaSRModel, SwinaSRForImageSuperResolution) if is_torch_available() else ()
      A__						:		Optional[Any]							      =					{'''feature-extraction''': SwinaSRModel} if is_torch_available() else {}

      A__						:		List[str]							      =					False
      A__						:		Optional[Any]							      =					False
      A__						:		str							      =					False
      A__						:		Optional[int]							      =					False
      A__						:		List[Any]							      =					False






      def       lowerCamelCase__     (     self						)						->Union[str, Any]:







        '''simple docstring'''


        _a					      =						SwinaSRModelTester(self						)
        _a					      =						ConfigTester(self  ,  config_class=a_  ,  embed_dim=3_7						)






      def       lowerCamelCase__     (     self						)						->Union[str, Any]:







        '''simple docstring'''


        self.config_tester.create_and_test_config_to_json_string()
        self.config_tester.create_and_test_config_to_json_file()
        self.config_tester.create_and_test_config_from_and_save_pretrained()
        self.config_tester.create_and_test_config_with_num_labels()
        self.config_tester.check_config_can_be_init_without_params()
        self.config_tester.check_config_arguments_init()






      def       lowerCamelCase__     (     self						)						->int:







        '''simple docstring'''


        _a					      =						self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*a_						)






      def       lowerCamelCase__     (     self						)						->Optional[int]:







        '''simple docstring'''


        _a					      =						self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_image_super_resolution(*a_						)






      @unittest.skip(reason="Got `CUDA error: misaligned address` with PyTorch 2.0.0."						)
      def       lowerCamelCase__     (     self						)						->int:







        '''simple docstring'''


        pass






      @unittest.skip(reason="Swin2SR does not use inputs_embeds"						)
      def       lowerCamelCase__     (     self						)						->Union[str, Any]:







        '''simple docstring'''


        pass






      @unittest.skip(reason="Swin2SR does not support training yet"						)
      def       lowerCamelCase__     (     self						)						->Union[str, Any]:







        '''simple docstring'''


        pass






      @unittest.skip(reason="Swin2SR does not support training yet"						)
      def       lowerCamelCase__     (     self						)						->Any:







        '''simple docstring'''


        pass






      def       lowerCamelCase__     (     self						)						->Tuple:







        '''simple docstring'''


        _a      ,						_a					      =						self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
          _a					      =						model_class(a_						)
          self.assertIsInstance(model.get_input_embeddings()  ,  (nn.Module)						)
          _a					      =						model.get_output_embeddings()
          self.assertTrue(x is None or isinstance(a_  ,  nn.Linear						)						)






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







        '''simple docstring'''


        _a      ,						_a					      =						self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
          _a					      =						model_class(a_						)
          _a					      =						inspect.signature(model.forward						)
          # signature.parameters is an OrderedDict => so arg_names order is deterministic
          _a					      =						[*signature.parameters.keys()]

          _a					      =						["pixel_values"]
          self.assertListEqual(arg_names[:1]  ,  a_						)






      @slow
      def       lowerCamelCase__     (     self						)						->Dict:







        '''simple docstring'''


        for model_name in SWIN2SR_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
          _a					      =						SwinaSRModel.from_pretrained(a_						)
          self.assertIsNotNone(a_						)






      def       lowerCamelCase__     (     self						)						->Dict:







        '''simple docstring'''


        _a      ,						_a					      =						self.model_tester.prepare_config_and_inputs_for_common()

        _a					      =						_config_zero_init(a_						)
        for model_class in self.all_model_classes:
          _a					      =						model_class(config=a_						)
          for name, param in model.named_parameters():
            if "logit_scale" in name:
              continue
            if param.requires_grad:
              self.assertIn(
                  ((param.data.mean() * 1E9).round() / 1E9).item()  ,  [0.0, 1.0]  ,  msg=f'''Parameter {name} of model {model_class} seems not properly initialized'''  ,  )






      def       lowerCamelCase__     (     self						)						->Tuple:







        '''simple docstring'''


        _a      ,						_a					      =						self.model_tester.prepare_config_and_inputs_for_common()
        _a					      =						True

        for model_class in self.all_model_classes:
          _a					      =						True
          _a					      =						False
          _a					      =						True
          _a					      =						model_class(a_						)
          model.to(a_						)
          model.eval()
          with torch.no_grad():
            _a					      =						model(**self._prepare_for_class(a_  ,  a_						)						)
          _a					      =						outputs.attentions
          _a					      =						len(self.model_tester.depths						)
          self.assertEqual(len(a_						)  ,  a_						)

          # check that output_attentions also work using config
          del inputs_dict["output_attentions"]
          _a					      =						True
          _a					      =						config.window_size**2
          _a					      =						model_class(a_						)
          model.to(a_						)
          model.eval()
          with torch.no_grad():
            _a					      =						model(**self._prepare_for_class(a_  ,  a_						)						)
          _a					      =						outputs.attentions
          self.assertEqual(len(a_						)  ,  a_						)

          self.assertListEqual(
              list(attentions[0].shape[-3:]						)  ,  [self.model_tester.num_heads[0], window_size_squared, window_size_squared]  ,  )
          _a					      =						len(a_						)

          # Check attention is always last and order is fine
          _a					      =						True
          _a					      =						True
          _a					      =						model_class(a_						)
          model.to(a_						)
          model.eval()
          with torch.no_grad():
            _a					      =						model(**self._prepare_for_class(a_  ,  a_						)						)

          self.assertEqual(out_len + 1  ,  len(a_						)						)

          _a					      =						outputs.attentions

          self.assertEqual(len(a_						)  ,  a_						)

          self.assertListEqual(
              list(self_attentions[0].shape[-3:]						)  ,  [self.model_tester.num_heads[0], window_size_squared, window_size_squared]  ,  )







@require_vision
@require_torch
@slow
class  lowercase_     (unittest.TestCase  ):






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







        '''simple docstring'''


        _a					      =						SwinaSRImageProcessor()
        _a					      =						SwinaSRForImageSuperResolution.from_pretrained("caidas/swin2SR-classical-sr-x2-64"						).to(a_						)

        _a					      =						Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"						)
        _a					      =						processor(images=a_  ,  return_tensors="pt"						).to(a_						)

        # forward pass
        with torch.no_grad():
          _a					      =						model(**a_						)

        # verify the logits
        _a					      =						torch.Size([1, 3, 9_7_6, 1_2_9_6]						)
        self.assertEqual(outputs.reconstruction.shape  ,  a_						)
        _a					      =						torch.tensor(
            [[0.5_458, 0.5_546, 0.5_638], [0.5_526, 0.5_565, 0.5_651], [0.5_396, 0.5_426, 0.5_621]]						).to(a_						)
        self.assertTrue(torch.allclose(outputs.reconstruction[0, 0, :3, :3]  ,  a_  ,  atol=1E-4						)						)


