



'''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_					(       SCREAMING_SNAKE_CASE_  ):
  """simple docstring"""




  def __init__(				self   :			str     )      ->      Dict:
    self.register_to_config()
    super().__init__()




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




  def 		_lowercase  (				self   :			Tuple    ,			_a   :			Dict     )      ->      Union[str, Any]:
    __lowerCamelCase : Optional[Any]					 =				[]
    __lowerCamelCase : int					 =				config_dict.copy()
    for key in config_dict.keys():
      if key.startswith('_'     ):
        temp_dict.pop(_a     )
        meta_keys.append(_a     )
    return (temp_dict, meta_keys)




  @torch.no_grad()
  def 		_lowercase  (				self   :			Optional[int]    ,			_a   :			List[Union[str, os.PathLike]]    ,			**_a   :			str     )      ->      List[Any]:
    __lowerCamelCase : Optional[int]					 =				kwargs.pop('cache_dir'    ,			_a     )
    __lowerCamelCase : Optional[int]					 =				kwargs.pop('resume_download'    ,			_a     )
    __lowerCamelCase : Tuple					 =				kwargs.pop('force_download'    ,			_a     )
    __lowerCamelCase : Optional[Any]					 =				kwargs.pop('proxies'    ,			_a     )
    __lowerCamelCase : List[Any]					 =				kwargs.pop('local_files_only'    ,			_a     )
    __lowerCamelCase : Any					 =				kwargs.pop('use_auth_token'    ,			_a     )
    __lowerCamelCase : str					 =				kwargs.pop('revision'    ,			_a     )
    __lowerCamelCase : List[Any]					 =				kwargs.pop('torch_dtype'    ,			_a     )
    __lowerCamelCase : int					 =				kwargs.pop('device_map'    ,			_a     )

    __lowerCamelCase : int					 =				kwargs.pop('alpha'    ,			0.5     )
    __lowerCamelCase : List[Any]					 =				kwargs.pop('interp'    ,			_a     )

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

    __lowerCamelCase : Union[str, Any]					 =				len(_a     )
    # Ignore result from model_index_json comparision of the two checkpoints
    __lowerCamelCase : Optional[Any]					 =				kwargs.pop('force'    ,			_a     )

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

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

      __lowerCamelCase : Optional[Any]					 =				(
          pretrained_model_name_or_path
          if os.path.isdir(_a     )
          else snapshot_download(
              _a    ,			cache_dir=_a    ,			resume_download=_a    ,			proxies=_a    ,			local_files_only=_a    ,			use_auth_token=_a    ,			revision=_a    ,			allow_patterns=_a    ,			user_agent=_a    ,			)
      )
      print('Cached Folder'    ,			_a     )
      cached_folders.append(_a     )

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

    __lowerCamelCase : Tuple					 =				None
    if len(_a     ) > 2:
      __lowerCamelCase : Dict					 =				os.path.join(cached_folders[2]     )

    if interp == "sigmoid":
      __lowerCamelCase : Optional[Any]					 =				CheckpointMergerPipeline.sigmoid
    elif interp == "inv_sigmoid":
      __lowerCamelCase : Any					 =				CheckpointMergerPipeline.inv_sigmoid
    elif interp == "add_diff":
      __lowerCamelCase : str					 =				CheckpointMergerPipeline.add_difference
    else:
      __lowerCamelCase : Union[str, Any]					 =				CheckpointMergerPipeline.weighted_sum

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

          __lowerCamelCase : Optional[int]					 =				getattr(_a    ,			'load_state_dict'     )
          __lowerCamelCase : Dict					 =				(
              safetensors.torch.load_file(_a     )
              if (is_safetensors_available() and checkpoint_path_a.endswith('.safetensors'     ))
              else torch.load(_a    ,			map_location='cpu'     )
          )
          __lowerCamelCase : Any					 =				None
          if checkpoint_path_a:
            __lowerCamelCase : List[str]					 =				(
                safetensors.torch.load_file(_a     )
                if (is_safetensors_available() and checkpoint_path_a.endswith('.safetensors'     ))
                else torch.load(_a    ,			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(_a     )}'     )
          continue
        print(f'MERGING {attr}'     )

        for key in theta_a.keys():
          if theta_a:
            __lowerCamelCase : Tuple					 =				theta_func(theta_a[key]    ,			theta_a[key]    ,			theta_a[key]    ,			_a     )
          else:
            __lowerCamelCase : List[Any]					 =				theta_func(theta_a[key]    ,			theta_a[key]    ,			_a    ,			_a     )

        del theta_a
        del theta_a
        update_theta_a(_a     )

        del theta_a
    return final_pipe




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




  @staticmethod
  def 		_lowercase  (				_a   :			str    ,			_a   :			Optional[int]    ,			_a   :			int    ,			_a   :			Optional[int]     )      ->      Dict:
    __lowerCamelCase : Optional[Any]					 =				alpha * alpha * (3 - (2 * alpha))
    return thetaa + ((thetaa - thetaa) * alpha)




  @staticmethod
  def 		_lowercase  (				_a   :			Any    ,			_a   :			str    ,			_a   :			List[str]    ,			_a   :			Any     )      ->      Optional[Any]:
    import math

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




  @staticmethod
  def 		_lowercase  (				_a   :			int    ,			_a   :			int    ,			_a   :			Any    ,			_a   :			List[Any]     )      ->      Optional[Any]:
    return thetaa + (thetaa - thetaa) * (1.0 - alpha)











'''simple docstring'''





import json
import os
import subprocess
import sys
import warnings
from argparse import ArgumentParser
from contextlib import AbstractContextManager
from typing import Dict, List, Optional

import requests

from ..utils import logging
from . import BaseTransformersCLICommand


_UpperCamelCase							       =		logging.get_logger(__name__)  # pylint: disable=invalid-name


_UpperCamelCase							       =		'lfs-multipart-upload'






class       lowerCamelCase_					(       SCREAMING_SNAKE_CASE_  ):
  """simple docstring"""




  @staticmethod
  def 		_lowercase  (				_a   :			ArgumentParser     )      ->      List[Any]:
    __lowerCamelCase : Union[str, Any]					 =				parser.add_parser(
        'lfs-enable-largefiles'    ,			help=(
            'Deprecated: use `huggingface-cli` instead. Configure your repository to enable upload of files > 5GB.'
        )    ,			)
    enable_parser.add_argument('path'    ,			type=_a    ,			help='Local path to repository you want to configure.'     )
    enable_parser.set_defaults(func=lambda _a     : LfsEnableCommand(_a     )     )

    __lowerCamelCase : Any					 =				parser.add_parser(
        _a    ,			help=(
            'Deprecated: use `huggingface-cli` instead. '
            'Command will get called by git-lfs, do not call it directly.'
        )    ,			)
    upload_parser.set_defaults(func=lambda _a     : LfsUploadCommand(_a     )     )






class       lowerCamelCase_					:
  """simple docstring"""




  def __init__(				self   :			Union[str, Any]    ,			_a   :			Any     )      ->      Tuple:
    __lowerCamelCase : int					 =				args




  def 		_lowercase  (				self   :			Union[str, Any]     )      ->      List[str]:
    warnings.warn(
        'Managing repositories through transformers-cli is deprecated. Please use `huggingface-cli` instead.'     )
    __lowerCamelCase : str					 =				os.path.abspath(self.args.path     )
    if not os.path.isdir(_a     ):
      print('This does not look like a valid git repo.'     )
      exit(1     )
    subprocess.run(
        'git config lfs.customtransfer.multipart.path transformers-cli'.split()    ,			check=_a    ,			cwd=_a     )
    subprocess.run(
        f'git config lfs.customtransfer.multipart.args {LFS_MULTIPART_UPLOAD_COMMAND}'.split()    ,			check=_a    ,			cwd=_a    ,			)
    print('Local repo set up for largefiles'     )




def  a_  (							_lowerCAmelCase )   ->   List[str]:
  __lowerCamelCase : int					 =				json.dumps(_lowerCAmelCase ) + '\n'
  sys.stdout.write(_lowerCAmelCase )
  sys.stdout.flush()

def  a_  (							)   ->   Optional[Dict]:
  __lowerCamelCase : str					 =				json.loads(sys.stdin.readline().strip() )

  if "terminate" in (msg.get('type' ), msg.get('event' )):
    # terminate message received
    return None

  if msg.get('event' ) not in ("download", "upload"):
    logger.critical('Received unexpected message' )
    sys.exit(1 )

  return msg






class       lowerCamelCase_					(       SCREAMING_SNAKE_CASE_  ):
  """simple docstring"""




  def __init__(				self   :			int    ,			_a   :			str    ,			_a   :			int    ,			_a   :			int     )      ->      Optional[Any]:
    __lowerCamelCase : Any					 =				filepath
    __lowerCamelCase : Dict					 =				seek_from
    __lowerCamelCase : Dict					 =				read_limit
    __lowerCamelCase : Tuple					 =				0




  def __enter__(				self   :			int     )      ->      int:
    __lowerCamelCase : Tuple					 =				open(self.filepath    ,			'rb'     )
    self.f.seek(self.seek_from     )
    return self




  def __len__(				self   :			Optional[int]     )      ->      Any:
    __lowerCamelCase : int					 =				os.fstat(self.f.fileno()     ).st_size
    return min(self.read_limit    ,			total_length - self.seek_from     )




  def 		_lowercase  (				self   :			Tuple    ,			_a   :			Optional[int]=-1     )      ->      Optional[int]:
    if self.n_seen >= self.read_limit:
      return b""
    __lowerCamelCase : Union[str, Any]					 =				self.read_limit - self.n_seen
    __lowerCamelCase : Tuple					 =				self.f.read(remaining_amount if n < 0 else min(_a    ,			_a     )     )
    self.n_seen += len(_a     )
    return data




  def __iter__(				self   :			int     )      ->      List[str]:
    yield self.read(n=4 * 1024 * 1024     )




  def __exit__(				self   :			List[str]    ,			*_a   :			int     )      ->      Optional[Any]:
    self.f.close()






class       lowerCamelCase_					:
  """simple docstring"""




  def __init__(				self   :			Optional[Any]    ,			_a   :			List[Any]     )      ->      Union[str, Any]:
    __lowerCamelCase : int					 =				args




  def 		_lowercase  (				self   :			Union[str, Any]     )      ->      Tuple:
    # Immediately after invoking a custom transfer process, git-lfs
    # sends initiation data to the process over stdin.
    # This tells the process useful information about the configuration.
    __lowerCamelCase : List[str]					 =				json.loads(sys.stdin.readline().strip()     )
    if not (init_msg.get('event'     ) == "init" and init_msg.get('operation'     ) == "upload"):
      write_msg({'error': {'code': 32, 'message': 'Wrong lfs init operation'}}     )
      sys.exit(1     )

    # The transfer process should use the information it needs from the
    # initiation structure, and also perform any one-off setup tasks it
    # needs to do. It should then respond on stdout with a simple empty
    # confirmation structure, as follows:
    write_msg({}     )

    # After the initiation exchange, git-lfs will send any number of
    # transfer requests to the stdin of the transfer process, in a serial sequence.
    while True:
      __lowerCamelCase : Any					 =				read_msg()
      if msg is None:
        # When all transfers have been processed, git-lfs will send
        # a terminate event to the stdin of the transfer process.
        # On receiving this message the transfer process should
        # clean up and terminate. No response is expected.
        sys.exit(0     )

      __lowerCamelCase : Tuple					 =				msg['oid']
      __lowerCamelCase : Dict					 =				msg['path']
      __lowerCamelCase : Any					 =				msg['action']['href']
      __lowerCamelCase : List[Any]					 =				msg['action']['header']
      __lowerCamelCase : List[str]					 =				int(header.pop('chunk_size'     )     )
      __lowerCamelCase : List[str]					 =				list(header.values()     )

      __lowerCamelCase : Any					 =				[]
      for i, presigned_url in enumerate(_a     ):
        with FileSlice(_a    ,			seek_from=i * chunk_size    ,			read_limit=_a     ) as data:
          __lowerCamelCase : Union[str, Any]					 =				requests.put(_a    ,			data=_a     )
          r.raise_for_status()
          parts.append(
              {
                  'etag': r.headers.get('etag'     ),
                  'partNumber': i + 1,
              }     )
          # In order to support progress reporting while data is uploading / downloading,
          # the transfer process should post messages to stdout
          write_msg(
              {
                  'event': 'progress',
                  'oid': oid,
                  'bytesSoFar': (i + 1) * chunk_size,
                  'bytesSinceLast': chunk_size,
              }     )
          # Not precise but that's ok.

      __lowerCamelCase : Tuple					 =				requests.post(
          _a    ,			json={
              'oid': oid,
              'parts': parts,
          }    ,			)
      r.raise_for_status()

      write_msg({'event': 'complete', 'oid': oid}     )






