





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







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



        '''simple docstring'''




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







      def      _lowercase						(					self	,   lowerCAmelCase__	,   lowerCAmelCase__							)     ->		Any:



        '''simple docstring'''




        if dicta == dicta:
          return True
        else:
          a__			:      List[Any]            =self._remove_meta_keys(__lowerCamelCase							)
          a__			:      Tuple            =self._remove_meta_keys(__lowerCamelCase							)
          if configa == configa:
            print(F'''Warning !: Mismatch in keys {meta_keysa} and {meta_keysa}.'''							)
            return True
        return False







      def      _lowercase						(					self	,   lowerCAmelCase__							)     ->		Tuple:



        '''simple docstring'''




        a__			:      Any            =[]
        a__			:      List[Any]            =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      _lowercase						(					self	,   lowerCAmelCase__	,   **lowerCAmelCase__							)     ->		List[Any]:



        '''simple docstring'''




        a__			:      Any            =kwargs.pop("cache_dir"	,   __lowerCamelCase							)
        a__			:      List[Any]            =kwargs.pop("resume_download"	,   __lowerCamelCase							)
        a__			:      int            =kwargs.pop("force_download"	,   __lowerCamelCase							)
        a__			:      List[Any]            =kwargs.pop("proxies"	,   __lowerCamelCase							)
        a__			:      Optional[int]            =kwargs.pop("local_files_only"	,   __lowerCamelCase							)
        a__			:      Union[str, Any]            =kwargs.pop("use_auth_token"	,   __lowerCamelCase							)
        a__			:      int            =kwargs.pop("revision"	,   __lowerCamelCase							)
        a__			:      Tuple            =kwargs.pop("torch_dtype"	,   __lowerCamelCase							)
        a__			:      List[Any]            =kwargs.pop("device_map"	,   __lowerCamelCase							)

        a__			:      List[Any]            =kwargs.pop("alpha"	,   0.5							)
        a__			:      Any            =kwargs.pop("interp"	,   __lowerCamelCase							)

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

        a__			:      Tuple            =len(__lowerCamelCase							)
        # Ignore result from model_index_json comparision of the two checkpoints
        a__			:      Tuple            =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__			:      Union[str, Any]            =[]
        for pretrained_model_name_or_path in pretrained_model_name_or_path_list:
          a__			:      Dict            =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__			:      Any            =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__			:      Optional[Any]            =[]
        for pretrained_model_name_or_path, config_dict in zip(__lowerCamelCase	,   __lowerCamelCase							):
          a__			:      int            =[k for k in config_dict.keys() if not k.startswith("_"							)]
          a__			:      Any            =[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__			:      int            =config_dict.get("_class_name"							)
          a__			:      int            ={"diffusers": __version__, "pipeline_class": requested_pipeline_class}

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

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

        if interp == "sigmoid":
          a__			:      List[str]            =CheckpointMergerPipeline.sigmoid
        elif interp == "inv_sigmoid":
          a__			:      Dict            =CheckpointMergerPipeline.inv_sigmoid
        elif interp == "add_diff":
          a__			:      Dict            =CheckpointMergerPipeline.add_difference
        else:
          a__			:      List[str]            =CheckpointMergerPipeline.weighted_sum

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

              a__			:      List[str]            =getattr(__lowerCamelCase	,   "load_state_dict"							)
              a__			:      List[Any]            =(
                  safetensors.torch.load_file(__lowerCamelCase							)
                  if (is_safetensors_available() and checkpoint_path_a.endswith(".safetensors"							))
                  else torch.load(__lowerCamelCase	,   map_location="cpu"							)
              )
              a__			:      List[str]            =None
              if checkpoint_path_a:
                a__			:      Union[str, Any]            =(
                    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__			:      List[str]            =theta_func(theta_a[key]	,   theta_a[key]	,   theta_a[key]	,   __lowerCamelCase							)
              else:
                a__			:      Optional[int]            =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      _lowercase						(					lowerCAmelCase__	,   lowerCAmelCase__	,   lowerCAmelCase__	,   lowerCAmelCase__							)     ->		str:



        '''simple docstring'''




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







      @staticmethod
      def      _lowercase						(					lowerCAmelCase__	,   lowerCAmelCase__	,   lowerCAmelCase__	,   lowerCAmelCase__							)     ->		List[Any]:



        '''simple docstring'''




        a__			:      Tuple            =alpha * alpha * (3 - (2 * alpha))
        return thetaa + ((thetaa - thetaa) * alpha)







      @staticmethod
      def      _lowercase						(					lowerCAmelCase__	,   lowerCAmelCase__	,   lowerCAmelCase__	,   lowerCAmelCase__							)     ->		Union[str, Any]:



        '''simple docstring'''




        import math

        a__			:      Optional[int]            =0.5 - math.sin(math.asin(1.0 - 2.0 * alpha							) / 3.0							)
        return thetaa + ((thetaa - thetaa) * alpha)







      @staticmethod
      def      _lowercase						(					lowerCAmelCase__	,   lowerCAmelCase__	,   lowerCAmelCase__	,   lowerCAmelCase__							)     ->		Optional[Any]:



        '''simple docstring'''




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













import os
import tempfile
from pathlib import Path
from unittest import TestCase

import pyarrow as pa
import pytest

from datasets.arrow_dataset import Dataset
from datasets.arrow_reader import ArrowReader, BaseReader, FileInstructions, ReadInstruction, make_file_instructions
from datasets.info import DatasetInfo
from datasets.splits import NamedSplit, Split, SplitDict, SplitInfo

from .utils import assert_arrow_memory_doesnt_increase, assert_arrow_memory_increases







class        __lowerCAmelCase (					UpperCamelCase__):







      def      _lowercase						(					self	,   lowerCAmelCase__	,   lowerCAmelCase__=False							)     ->		List[Any]:



        '''simple docstring'''




        a__      ,							a__      ,							a__			:      Optional[int]            =(
            filename_skip_take["filename"],
            filename_skip_take["skip"] if "skip" in filename_skip_take else None,
            filename_skip_take["take"] if "take" in filename_skip_take else None,
        )
        open(os.path.join(lowerCAmelCase__							)	,   "wb"							).close()
        a__			:      Optional[int]            =pa.Table.from_pydict({"filename": [Path(lowerCAmelCase__							).name] * 1_0_0}							)
        if take == -1:
          a__			:      int            =len(lowerCAmelCase__							) - skip
        if skip is not None and take is not None:
          a__			:      Optional[Any]            =pa_table.slice(lowerCAmelCase__	,   lowerCAmelCase__							)
        return pa_table







class        __lowerCAmelCase (					UpperCamelCase__):







      def      _lowercase						(					self							)     ->		str:



        '''simple docstring'''




        a__			:      List[str]            ="my_name"
        a__			:      List[Any]            =SplitInfo(name="train"	,   num_examples=1_0_0							)
        a__			:      List[str]            =SplitInfo(name="test"	,   num_examples=1_0_0							)
        a__			:      List[Any]            =[train_info, test_info]
        a__			:      Union[str, Any]            =SplitDict()
        split_dict.add(lowerCAmelCase__							)
        split_dict.add(lowerCAmelCase__							)
        a__			:      Dict            =DatasetInfo(splits=lowerCAmelCase__							)

        with tempfile.TemporaryDirectory() as tmp_dir:
          a__			:      Optional[Any]            =ReaderTest(lowerCAmelCase__	,   lowerCAmelCase__							)

          a__			:      Any            ="test[:33%]"
          a__			:      int            =Dataset(**reader.read(lowerCAmelCase__	,   lowerCAmelCase__	,   lowerCAmelCase__							)							)
          self.assertEqual(dset["filename"][0]	,   F'''{name}-test'''							)
          self.assertEqual(dset.num_rows	,   3_3							)
          self.assertEqual(dset.num_columns	,   1							)

          a__			:      Tuple            =["train", "test[:33%]"]
          a__			:      str            =[Split.TRAIN, ReadInstruction.from_spec("test[:33%]"							)]
          for instructions in [instructionsa, instructionsa]:
            a__			:      Optional[int]            =[reader.read(lowerCAmelCase__	,   lowerCAmelCase__	,   lowerCAmelCase__							) for instr in instructions]
            a__      ,							a__			:      Optional[int]            =(Dataset(**lowerCAmelCase__							) for dataset_kwargs in datasets_kwargs)
            self.assertEqual(train_dset["filename"][0]	,   F'''{name}-train'''							)
            self.assertEqual(train_dset.num_rows	,   1_0_0							)
            self.assertEqual(train_dset.num_columns	,   1							)
            self.assertIsInstance(train_dset.split	,   lowerCAmelCase__							)
            self.assertEqual(str(train_dset.split							)	,   "train"							)
            self.assertEqual(test_dset["filename"][0]	,   F'''{name}-test'''							)
            self.assertEqual(test_dset.num_rows	,   3_3							)
            self.assertEqual(test_dset.num_columns	,   1							)
            self.assertIsInstance(test_dset.split	,   lowerCAmelCase__							)
            self.assertEqual(str(test_dset.split							)	,   "test[:33%]"							)
            del train_dset, test_dset







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



        '''simple docstring'''




        a__			:      Any            ="my_name"
        a__			:      List[Any]            =SplitInfo(name="train"	,   num_examples=1_0_0_0	,   shard_lengths=[1_0_0] * 1_0							)
        a__			:      Optional[Any]            =[train_info]
        a__			:      Any            =SplitDict()
        split_dict.add(lowerCAmelCase__							)
        a__			:      Any            =DatasetInfo(splits=lowerCAmelCase__							)

        with tempfile.TemporaryDirectory() as tmp_dir:
          a__			:      Optional[int]            =ReaderTest(lowerCAmelCase__	,   lowerCAmelCase__							)

          a__			:      Union[str, Any]            ="train[:33%]"
          a__			:      Union[str, Any]            =Dataset(**reader.read(lowerCAmelCase__	,   lowerCAmelCase__	,   lowerCAmelCase__							)							)
          self.assertEqual(dset["filename"][0]	,   F'''{name}-train-00000-of-00010'''							)
          self.assertEqual(dset["filename"][-1]	,   F'''{name}-train-00003-of-00010'''							)
          self.assertEqual(dset.num_rows	,   3_3_0							)
          self.assertEqual(dset.num_columns	,   1							)







      def      _lowercase						(					self							)     ->		List[str]:



        '''simple docstring'''




        a__			:      Dict            =SplitInfo(name="train"	,   num_examples=1_0_0							)
        a__			:      str            =SplitInfo(name="test"	,   num_examples=1_0_0							)
        a__			:      int            =SplitDict()
        split_dict.add(lowerCAmelCase__							)
        split_dict.add(lowerCAmelCase__							)
        a__			:      int            =DatasetInfo(splits=lowerCAmelCase__							)

        with tempfile.TemporaryDirectory() as tmp_dir:
          a__			:      Tuple            =ReaderTest(lowerCAmelCase__	,   lowerCAmelCase__							)

          a__			:      Tuple            =[
              {"filename": os.path.join(lowerCAmelCase__	,   "train"							)},
              {"filename": os.path.join(lowerCAmelCase__	,   "test"							), "skip": 1_0, "take": 1_0},
          ]
          a__			:      Optional[Any]            =Dataset(**reader.read_files(lowerCAmelCase__	,   original_instructions="train+test[10:20]"							)							)
          self.assertEqual(dset.num_rows	,   1_1_0							)
          self.assertEqual(dset.num_columns	,   1							)
          del dset
@pytest.mark.parametrize("in_memory" , [False, True]		)
def   _A (	SCREAMING_SNAKE_CASE  :     Optional[Any] , SCREAMING_SNAKE_CASE  :     Optional[Any] , SCREAMING_SNAKE_CASE  :     Optional[Any]		):





  """simple docstring"""



  a__			:      int            =arrow_file
  with assert_arrow_memory_increases() if in_memory else assert_arrow_memory_doesnt_increase():
    a__			:      str            =ArrowReader.read_table(SCREAMING_SNAKE_CASE , in_memory=SCREAMING_SNAKE_CASE		)
  assert table.shape == dataset.data.shape
  assert set(table.column_names		) == set(dataset.data.column_names		)
  assert dict(table.to_pydict()		) == dict(dataset.data.to_pydict()		)  # to_pydict returns OrderedDict
@pytest.mark.parametrize("in_memory" , [False, True]		)
def   _A (	SCREAMING_SNAKE_CASE  :     Tuple , SCREAMING_SNAKE_CASE  :     List[Any] , SCREAMING_SNAKE_CASE  :     str		):





  """simple docstring"""



  a__			:      Optional[int]            =arrow_file
  a__			:      Tuple            =ArrowReader("" , SCREAMING_SNAKE_CASE		)
  with assert_arrow_memory_increases() if in_memory else assert_arrow_memory_doesnt_increase():
    a__			:      Union[str, Any]            =reader.read_files([{"filename": filename}] , in_memory=SCREAMING_SNAKE_CASE		)
  assert dataset_kwargs.keys() == {"arrow_table", "info", "split"}
  a__			:      Dict            =dataset_kwargs["arrow_table"]
  assert table.shape == dataset.data.shape
  assert set(table.column_names		) == set(dataset.data.column_names		)
  assert dict(table.to_pydict()		) == dict(dataset.data.to_pydict()		)  # to_pydict returns OrderedDict
def   _A (	):





  """simple docstring"""



  assert ReadInstruction("train" , to=10 , unit="abs"		).to_spec() == "train[:10]"
  assert ReadInstruction("train" , from_=-80 , to=10 , unit="%"		).to_spec() == "train[-80%:10%]"

  a__			:      List[Any]            ="train+test"
  assert ReadInstruction.from_spec(SCREAMING_SNAKE_CASE		).to_spec() == spec_train_test

  a__			:      Union[str, Any]            ="train[2:10]"
  assert ReadInstruction.from_spec(SCREAMING_SNAKE_CASE		).to_spec() == spec_train_abs

  a__			:      List[Any]            ="train[15%:-20%]"
  assert ReadInstruction.from_spec(SCREAMING_SNAKE_CASE		).to_spec() == spec_train_pct

  a__			:      Dict            ="train[:10%](closest)"
  assert ReadInstruction.from_spec(SCREAMING_SNAKE_CASE		).to_spec() == "train[:10%]"

  a__			:      Tuple            ="train[:10%](pct1_dropremainder)"
  assert ReadInstruction.from_spec(SCREAMING_SNAKE_CASE		).to_spec() == spec_train_pct_rounding

  a__			:      List[Any]            ="train[:10%](pct1_dropremainder)+test[-10%:](pct1_dropremainder)"
  assert ReadInstruction.from_spec(SCREAMING_SNAKE_CASE		).to_spec() == spec_train_test_pct_rounding
def   _A (	):





  """simple docstring"""



  a__			:      Optional[Any]            ="dummy"
  a__			:      Dict            =[SplitInfo(name="train" , num_examples=100		)]
  a__			:      Optional[Any]            ="train[:33%]"
  a__			:      int            ="arrow"
  a__			:      Any            ="prefix"

  a__			:      Tuple            =make_file_instructions(SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE		)
  assert isinstance(SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE		)
  assert file_instructions.num_examples == 33
  assert file_instructions.file_instructions == [
      {"filename": os.path.join(SCREAMING_SNAKE_CASE , f'''{name}-train.arrow'''		), "skip": 0, "take": 33}
  ]

  a__			:      List[str]            =[SplitInfo(name="train" , num_examples=100 , shard_lengths=[10] * 10		)]
  a__			:      Union[str, Any]            =make_file_instructions(SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE		)
  assert isinstance(SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE		)
  assert file_instructions.num_examples == 33
  assert file_instructions.file_instructions == [
      {"filename": os.path.join(SCREAMING_SNAKE_CASE , f'''{name}-train-00000-of-00010.arrow'''		), "skip": 0, "take": -1},
      {"filename": os.path.join(SCREAMING_SNAKE_CASE , f'''{name}-train-00001-of-00010.arrow'''		), "skip": 0, "take": -1},
      {"filename": os.path.join(SCREAMING_SNAKE_CASE , f'''{name}-train-00002-of-00010.arrow'''		), "skip": 0, "take": -1},
      {"filename": os.path.join(SCREAMING_SNAKE_CASE , f'''{name}-train-00003-of-00010.arrow'''		), "skip": 0, "take": 3},
  ]





@pytest.mark.parametrize("name, expected_exception" , [(None, TypeError), ("", ValueError)]		)
def   _A (	SCREAMING_SNAKE_CASE  :     Tuple , SCREAMING_SNAKE_CASE  :     Tuple		):





  """simple docstring"""



  a__			:      Dict            =[SplitInfo(name="train" , num_examples=100		)]
  a__			:      Any            ="train"
  a__			:      Optional[Any]            ="arrow"
  a__			:      List[Any]            ="prefix_path"
  with pytest.raises(SCREAMING_SNAKE_CASE		):
    a__			:      Optional[Any]            =make_file_instructions(SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE		)






