



'''simple docstring'''





import inspect
import tempfile
from collections import OrderedDict, UserDict
from collections.abc import MutableMapping
from contextlib import ExitStack, contextmanager
from dataclasses import fields
from enum import Enum
from typing import Any, ContextManager, List, Tuple

import numpy as np

from .import_utils import is_flax_available, is_tf_available, is_torch_available, is_torch_fx_proxy


if is_flax_available():
       import jax.numpy as jnp







class     UpperCAmelCase_       (       _a	):
    """simple docstring"""







    def __get__(    self					:     str ,     snake_case_					:     Union[str, Any] ,     snake_case_					:     Tuple=None						):
          # See docs.python.org/3/howto/descriptor.html#properties
          if obj is None:
                return self
          if self.fget is None:
                raise AttributeError("""unreadable attribute"""						)
          snake_case__  :    Optional[int]    						=      """__cached_""" + self.fget.__name__
          snake_case__  :    int    						=      getattr(snake_case_ ,     snake_case_ ,     snake_case_						)
          if cached is None:
                snake_case__  :    Union[str, Any]    						=      self.fget(snake_case_						)
                setattr(snake_case_ ,     snake_case_ ,     snake_case_						)
          return cached





def 						__snake_case(				_lowerCAmelCase		)		-> Tuple:
      snake_case__  :    Optional[int]    						=      val.lower()
      if val in {"y", "yes", "t", "true", "on", "1"}:
            return 1
      if val in {"n", "no", "f", "false", "off", "0"}:
            return 0
      raise ValueError(f"invalid truth value {val!r}"		)





def 						__snake_case(				_lowerCAmelCase		)		-> Any:
      if is_torch_fx_proxy(_lowerCAmelCase		):
            return True
      if is_torch_available():
            import torch

            if isinstance(_lowerCAmelCase     ,   torch.Tensor		):
                  return True
      if is_tf_available():
            import tensorflow as tf

            if isinstance(_lowerCAmelCase     ,   tf.Tensor		):
                  return True

      if is_flax_available():
            import jax.numpy as jnp
            from jax.core import Tracer

            if isinstance(_lowerCAmelCase     ,   (jnp.ndarray, Tracer)		):
                  return True

      return isinstance(_lowerCAmelCase     ,   np.ndarray		)





def 						__snake_case(				_lowerCAmelCase		)		-> Optional[Any]:
      return isinstance(_lowerCAmelCase     ,   np.ndarray		)





def 						__snake_case(				_lowerCAmelCase		)		-> int:
      return _is_numpy(_lowerCAmelCase		)





def 						__snake_case(				_lowerCAmelCase		)		-> List[str]:
      import torch

      return isinstance(_lowerCAmelCase     ,   torch.Tensor		)





def 						__snake_case(				_lowerCAmelCase		)		-> str:
      return False if not is_torch_available() else _is_torch(_lowerCAmelCase		)





def 						__snake_case(				_lowerCAmelCase		)		-> List[str]:
      import torch

      return isinstance(_lowerCAmelCase     ,   torch.device		)





def 						__snake_case(				_lowerCAmelCase		)		-> List[Any]:
      return False if not is_torch_available() else _is_torch_device(_lowerCAmelCase		)





def 						__snake_case(				_lowerCAmelCase		)		-> List[Any]:
      import torch

      if isinstance(_lowerCAmelCase     ,   _lowerCAmelCase		):
            if hasattr(_lowerCAmelCase     ,   _lowerCAmelCase		):
                  snake_case__  :    Union[str, Any]    						=      getattr(_lowerCAmelCase     ,   _lowerCAmelCase		)
            else:
                  return False
      return isinstance(_lowerCAmelCase     ,   torch.dtype		)





def 						__snake_case(				_lowerCAmelCase		)		-> List[Any]:
      return False if not is_torch_available() else _is_torch_dtype(_lowerCAmelCase		)





def 						__snake_case(				_lowerCAmelCase		)		-> List[str]:
      import tensorflow as tf

      return isinstance(_lowerCAmelCase     ,   tf.Tensor		)





def 						__snake_case(				_lowerCAmelCase		)		-> Any:
      return False if not is_tf_available() else _is_tensorflow(_lowerCAmelCase		)





def 						__snake_case(				_lowerCAmelCase		)		-> Optional[int]:
      import tensorflow as tf

      # the `is_symbolic_tensor` predicate is only available starting with TF 2.14
      if hasattr(_lowerCAmelCase     ,   """is_symbolic_tensor"""		):
            return tf.is_symbolic_tensor(_lowerCAmelCase		)
      return type(_lowerCAmelCase		) == tf.Tensor





def 						__snake_case(				_lowerCAmelCase		)		-> Optional[Any]:
      return False if not is_tf_available() else _is_tf_symbolic_tensor(_lowerCAmelCase		)





def 						__snake_case(				_lowerCAmelCase		)		-> Dict:
      import jax.numpy as jnp  # noqa: F811

      return isinstance(_lowerCAmelCase     ,   jnp.ndarray		)





def 						__snake_case(				_lowerCAmelCase		)		-> List[str]:
      return False if not is_flax_available() else _is_jax(_lowerCAmelCase		)





def 						__snake_case(				_lowerCAmelCase		)		-> Any:
      if isinstance(_lowerCAmelCase     ,   (dict, UserDict)		):
            return {k: to_py_obj(_lowerCAmelCase		) for k, v in obj.items()}
      elif isinstance(_lowerCAmelCase     ,   (list, tuple)		):
            return [to_py_obj(_lowerCAmelCase		) for o in obj]
      elif is_tf_tensor(_lowerCAmelCase		):
            return obj.numpy().tolist()
      elif is_torch_tensor(_lowerCAmelCase		):
            return obj.detach().cpu().tolist()
      elif is_jax_tensor(_lowerCAmelCase		):
            return np.asarray(_lowerCAmelCase		).tolist()
      elif isinstance(_lowerCAmelCase     ,   (np.ndarray, np.number)		):  # tolist also works on 0d np arrays
            return obj.tolist()
      else:
            return obj





def 						__snake_case(				_lowerCAmelCase		)		-> List[Any]:
      if isinstance(_lowerCAmelCase     ,   (dict, UserDict)		):
            return {k: to_numpy(_lowerCAmelCase		) for k, v in obj.items()}
      elif isinstance(_lowerCAmelCase     ,   (list, tuple)		):
            return np.array(_lowerCAmelCase		)
      elif is_tf_tensor(_lowerCAmelCase		):
            return obj.numpy()
      elif is_torch_tensor(_lowerCAmelCase		):
            return obj.detach().cpu().numpy()
      elif is_jax_tensor(_lowerCAmelCase		):
            return np.asarray(_lowerCAmelCase		)
      else:
            return obj







class     UpperCAmelCase_       (       _a	):
    """simple docstring"""







    def lowerCamelCase       (    self					:     Union[str, Any]						):
          snake_case__  :    List[str]    						=      fields(self						)

          # Safety and consistency checks
          if not len(snake_case_						):
                raise ValueError(f"{self.__class__.__name__} has no fields."						)
          if not all(field.default is None for field in class_fields[1:]						):
                raise ValueError(f"{self.__class__.__name__} should not have more than one required field."						)

          snake_case__  :    List[str]    						=      getattr(self ,     class_fields[0].name						)
          snake_case__  :    int    						=      all(getattr(self ,     field.name						) is None for field in class_fields[1:]						)

          if other_fields_are_none and not is_tensor(snake_case_						):
                if isinstance(snake_case_ ,     snake_case_						):
                      snake_case__  :    List[str]    						=      first_field.items()
                      snake_case__  :    Optional[int]    						=      True
                else:
                      try:
                            snake_case__  :    Optional[Any]    						=      iter(snake_case_						)
                            snake_case__  :    Optional[int]    						=      True
                      except TypeError:
                            snake_case__  :    Dict    						=      False

            # if we provided an iterator as first field and the iterator is a (key, value) iterator
            # set the associated fields
                if first_field_iterator:
                      for idx, element in enumerate(snake_case_						):
                            if (
                                not isinstance(snake_case_ ,     (list, tuple)						)
                                or not len(snake_case_						) == 2
                                or not isinstance(element[0] ,     snake_case_						)
                            ):
                                  if idx == 0:
                                        # If we do not have an iterator of key/values, set it as attribute
                                        snake_case__  :    int    						=      first_field
                                  else:
                                        # If we have a mixed iterator, raise an error
                                        raise ValueError(
                                            f"Cannot set key/value for {element}. It needs to be a tuple (key, value)."						)
                                  break
                            setattr(self ,     element[0] ,     element[1]						)
                            if element[1] is not None:
                                  snake_case__  :    Union[str, Any]    						=      element[1]
                elif first_field is not None:
                      snake_case__  :    List[str]    						=      first_field
          else:
                for field in class_fields:
                      snake_case__  :    List[Any]    						=      getattr(self ,     field.name						)
                      if v is not None:
                            snake_case__  :    Dict    						=      v




    def __delitem__(    self					:     Union[str, Any] ,     *snake_case_					:     Tuple ,     **snake_case_					:     int						):
          raise Exception(f"You cannot use ``__delitem__`` on a {self.__class__.__name__} instance."						)




    def lowerCamelCase       (    self					:     List[str] ,     *snake_case_					:     List[Any] ,     **snake_case_					:     Optional[Any]						):
          raise Exception(f"You cannot use ``setdefault`` on a {self.__class__.__name__} instance."						)




    def lowerCamelCase       (    self					:     Dict ,     *snake_case_					:     List[Any] ,     **snake_case_					:     str						):
          raise Exception(f"You cannot use ``pop`` on a {self.__class__.__name__} instance."						)




    def lowerCamelCase       (    self					:     List[str] ,     *snake_case_					:     Tuple ,     **snake_case_					:     Tuple						):
          raise Exception(f"You cannot use ``update`` on a {self.__class__.__name__} instance."						)




    def __getitem__(    self					:     Tuple ,     snake_case_					:     Optional[Any]						):
          if isinstance(snake_case_ ,     snake_case_						):
                snake_case__  :    int    						=      dict(self.items()						)
                return inner_dict[k]
          else:
                return self.to_tuple()[k]




    def __setattr__(    self					:     List[str] ,     snake_case_					:     Tuple ,     snake_case_					:     Union[str, Any]						):
          if name in self.keys() and value is not None:
                # Don't call self.__setitem__ to avoid recursion errors
                super().__setitem__(snake_case_ ,     snake_case_						)
          super().__setattr__(snake_case_ ,     snake_case_						)




    def __setitem__(    self					:     str ,     snake_case_					:     Union[str, Any] ,     snake_case_					:     int						):
          # Will raise a KeyException if needed
          super().__setitem__(snake_case_ ,     snake_case_						)
          # Don't call self.__setattr__ to avoid recursion errors
          super().__setattr__(snake_case_ ,     snake_case_						)




    def lowerCamelCase       (    self					:     int						):
          return tuple(self[k] for k in self.keys()						)







class     UpperCAmelCase_       (       _a		,		_a	):
    """simple docstring"""







    @classmethod
    def lowerCamelCase       (    cls					:     int ,     snake_case_					:     List[Any]						):
          raise ValueError(
              f"{value} is not a valid {cls.__name__}, please select one of {list(cls._valueamember_map_.keys()						)}"						)







class     UpperCAmelCase_       (       _a	):
    """simple docstring"""







    lowercase					=		"longest"
    lowercase					=		"max_length"
    lowercase					=		"do_not_pad"







class     UpperCAmelCase_       (       _a	):
    """simple docstring"""







    lowercase					=		"pt"
    lowercase					=		"tf"
    lowercase					=		"np"
    lowercase					=		"jax"







class     UpperCAmelCase_       :
    """simple docstring"""







    def __init__(    self					:     int ,     snake_case_					:     List[ContextManager]						):
          snake_case__  :    Any    						=      context_managers
          snake_case__  :    int    						=      ExitStack()




    def __enter__(    self					:     Union[str, Any]						):
          for context_manager in self.context_managers:
                self.stack.enter_context(snake_case_						)




    def __exit__(    self					:     Optional[int] ,     *snake_case_					:     List[Any] ,     **snake_case_					:     Any						):
          self.stack.__exit__(*snake_case_ ,     **snake_case_						)





def 						__snake_case(				_lowerCAmelCase		)		-> List[str]:
      snake_case__  :    Optional[Any]    						=      infer_framework(_lowerCAmelCase		)
      if framework == "tf":
            snake_case__  :    Optional[Any]    						=      inspect.signature(model_class.call		)  # TensorFlow models
      elif framework == "pt":
            snake_case__  :    str    						=      inspect.signature(model_class.forward		)  # PyTorch models
      else:
            snake_case__  :    str    						=      inspect.signature(model_class.__call__		)  # Flax models

      for p in signature.parameters:
            if p == "return_loss" and signature.parameters[p].default is True:
                  return True

      return False





def 						__snake_case(				_lowerCAmelCase		)		-> Dict:
      snake_case__  :    Optional[int]    						=      model_class.__name__
      snake_case__  :    int    						=      infer_framework(_lowerCAmelCase		)
      if framework == "tf":
            snake_case__  :    List[str]    						=      inspect.signature(model_class.call		)  # TensorFlow models
      elif framework == "pt":
            snake_case__  :    int    						=      inspect.signature(model_class.forward		)  # PyTorch models
      else:
            snake_case__  :    Any    						=      inspect.signature(model_class.__call__		)  # Flax models

      if "QuestionAnswering" in model_name:
            return [p for p in signature.parameters if "label" in p or p in ("start_positions", "end_positions")]
      else:
            return [p for p in signature.parameters if "label" in p]





def 						__snake_case(				_lowerCAmelCase     ,   _lowerCAmelCase = ""     ,   _lowerCAmelCase = "."		)		-> List[str]:

      def _flatten_dict(_lowerCAmelCase     ,   _lowerCAmelCase=""     ,   _lowerCAmelCase="."		):
            for k, v in d.items():
                  snake_case__  :    Optional[int]    						=      str(_lowerCAmelCase		) + delimiter + str(_lowerCAmelCase		) if parent_key else k
                  if v and isinstance(_lowerCAmelCase     ,   _lowerCAmelCase		):
                        yield from flatten_dict(_lowerCAmelCase     ,   _lowerCAmelCase     ,   delimiter=_lowerCAmelCase		).items()
                  else:
                        yield key, v

      return dict(_flatten_dict(_lowerCAmelCase     ,   _lowerCAmelCase     ,   _lowerCAmelCase		)		)





@contextmanager
def 						__snake_case(				_lowerCAmelCase     ,   _lowerCAmelCase = False		)		-> Tuple:
      if use_temp_dir:
            with tempfile.TemporaryDirectory() as tmp_dir:
                  yield tmp_dir
      else:
            yield working_dir





def 						__snake_case(				_lowerCAmelCase     ,   _lowerCAmelCase=None		)		-> Optional[Any]:
      if is_numpy_array(_lowerCAmelCase		):
            return np.transpose(_lowerCAmelCase     ,   axes=_lowerCAmelCase		)
      elif is_torch_tensor(_lowerCAmelCase		):
            return array.T if axes is None else array.permute(*_lowerCAmelCase		)
      elif is_tf_tensor(_lowerCAmelCase		):
            import tensorflow as tf

            return tf.transpose(_lowerCAmelCase     ,   perm=_lowerCAmelCase		)
      elif is_jax_tensor(_lowerCAmelCase		):
            return jnp.transpose(_lowerCAmelCase     ,   axes=_lowerCAmelCase		)
      else:
            raise ValueError(f"Type not supported for transpose: {type(_lowerCAmelCase		)}."		)





def 						__snake_case(				_lowerCAmelCase     ,   _lowerCAmelCase		)		-> Tuple:
      if is_numpy_array(_lowerCAmelCase		):
            return np.reshape(_lowerCAmelCase     ,   _lowerCAmelCase		)
      elif is_torch_tensor(_lowerCAmelCase		):
            return array.reshape(*_lowerCAmelCase		)
      elif is_tf_tensor(_lowerCAmelCase		):
            import tensorflow as tf

            return tf.reshape(_lowerCAmelCase     ,   _lowerCAmelCase		)
      elif is_jax_tensor(_lowerCAmelCase		):
            return jnp.reshape(_lowerCAmelCase     ,   _lowerCAmelCase		)
      else:
            raise ValueError(f"Type not supported for reshape: {type(_lowerCAmelCase		)}."		)





def 						__snake_case(				_lowerCAmelCase     ,   _lowerCAmelCase=None		)		-> Dict:
      if is_numpy_array(_lowerCAmelCase		):
            return np.squeeze(_lowerCAmelCase     ,   axis=_lowerCAmelCase		)
      elif is_torch_tensor(_lowerCAmelCase		):
            return array.squeeze() if axis is None else array.squeeze(dim=_lowerCAmelCase		)
      elif is_tf_tensor(_lowerCAmelCase		):
            import tensorflow as tf

            return tf.squeeze(_lowerCAmelCase     ,   axis=_lowerCAmelCase		)
      elif is_jax_tensor(_lowerCAmelCase		):
            return jnp.squeeze(_lowerCAmelCase     ,   axis=_lowerCAmelCase		)
      else:
            raise ValueError(f"Type not supported for squeeze: {type(_lowerCAmelCase		)}."		)





def 						__snake_case(				_lowerCAmelCase     ,   _lowerCAmelCase		)		-> Tuple:
      if is_numpy_array(_lowerCAmelCase		):
            return np.expand_dims(_lowerCAmelCase     ,   _lowerCAmelCase		)
      elif is_torch_tensor(_lowerCAmelCase		):
            return array.unsqueeze(dim=_lowerCAmelCase		)
      elif is_tf_tensor(_lowerCAmelCase		):
            import tensorflow as tf

            return tf.expand_dims(_lowerCAmelCase     ,   axis=_lowerCAmelCase		)
      elif is_jax_tensor(_lowerCAmelCase		):
            return jnp.expand_dims(_lowerCAmelCase     ,   axis=_lowerCAmelCase		)
      else:
            raise ValueError(f"Type not supported for expand_dims: {type(_lowerCAmelCase		)}."		)





def 						__snake_case(				_lowerCAmelCase		)		-> Optional[Any]:
      if is_numpy_array(_lowerCAmelCase		):
            return np.size(_lowerCAmelCase		)
      elif is_torch_tensor(_lowerCAmelCase		):
            return array.numel()
      elif is_tf_tensor(_lowerCAmelCase		):
            import tensorflow as tf

            return tf.size(_lowerCAmelCase		)
      elif is_jax_tensor(_lowerCAmelCase		):
            return array.size
      else:
            raise ValueError(f"Type not supported for expand_dims: {type(_lowerCAmelCase		)}."		)





def 						__snake_case(				_lowerCAmelCase     ,   _lowerCAmelCase		)		-> Dict:
      for key, value in auto_map.items():
            if isinstance(_lowerCAmelCase     ,   (tuple, list)		):
                  snake_case__  :    Tuple    						=      [f"{repo_id}--{v}" if (v is not None and """--""" not in v) else v for v in value]
            elif value is not None and "--" not in value:
                  snake_case__  :    str    						=      f"{repo_id}--{value}"

      return auto_map





def 						__snake_case(				_lowerCAmelCase		)		-> int:
      for base_class in inspect.getmro(_lowerCAmelCase		):
            snake_case__  :    int    						=      base_class.__module__
            snake_case__  :    Any    						=      base_class.__name__
            if module.startswith("""tensorflow"""		) or module.startswith("""keras"""		) or name == "TFPreTrainedModel":
                  return "tf"
            elif module.startswith("""torch"""		) or name == "PreTrainedModel":
                  return "pt"
            elif module.startswith("""flax"""		) or module.startswith("""jax"""		) or name == "FlaxPreTrainedModel":
                  return "flax"
      else:
            raise TypeError(f"Could not infer framework from class {model_class}."		)

"""simple docstring"""






import argparse
import os.path as osp
import re

import torch
from safetensors.torch import load_file, save_file


# =================#
# UNet Conversion #
# =================#

lowerCAmelCase__  :  Optional[int]						     =	[
    # (stable-diffusion, HF Diffusers)
    ('time_embed.0.weight', 'time_embedding.linear_1.weight'),
    ('time_embed.0.bias', 'time_embedding.linear_1.bias'),
    ('time_embed.2.weight', 'time_embedding.linear_2.weight'),
    ('time_embed.2.bias', 'time_embedding.linear_2.bias'),
    ('input_blocks.0.0.weight', 'conv_in.weight'),
    ('input_blocks.0.0.bias', 'conv_in.bias'),
    ('out.0.weight', 'conv_norm_out.weight'),
    ('out.0.bias', 'conv_norm_out.bias'),
    ('out.2.weight', 'conv_out.weight'),
    ('out.2.bias', 'conv_out.bias'),
]

lowerCAmelCase__  :  str						     =	[
    # (stable-diffusion, HF Diffusers)
    ('in_layers.0', 'norm1'),
    ('in_layers.2', 'conv1'),
    ('out_layers.0', 'norm2'),
    ('out_layers.3', 'conv2'),
    ('emb_layers.1', 'time_emb_proj'),
    ('skip_connection', 'conv_shortcut'),
]

lowerCAmelCase__  :  int						     =	[]
# hardcoded number of downblocks and resnets/attentions...
# would need smarter logic for other networks.
for i in range(4):
						# loop over downblocks/upblocks

						for j in range(2):
												# loop over resnets/attentions for downblocks
												lowerCAmelCase__  :  Dict						     =	F"""down_blocks.{i}.resnets.{j}."""
												lowerCAmelCase__  :  Union[str, Any]						     =	F"""input_blocks.{3*i + j + 1}.0."""
												unet_conversion_map_layer.append((sd_down_res_prefix, hf_down_res_prefix))

												if i < 3:
																		# no attention layers in down_blocks.3
																		lowerCAmelCase__  :  Tuple						     =	F"""down_blocks.{i}.attentions.{j}."""
																		lowerCAmelCase__  :  int						     =	F"""input_blocks.{3*i + j + 1}.1."""
																		unet_conversion_map_layer.append((sd_down_atn_prefix, hf_down_atn_prefix))

						for j in range(3):
												# loop over resnets/attentions for upblocks
												lowerCAmelCase__  :  Tuple						     =	F"""up_blocks.{i}.resnets.{j}."""
												lowerCAmelCase__  :  int						     =	F"""output_blocks.{3*i + j}.0."""
												unet_conversion_map_layer.append((sd_up_res_prefix, hf_up_res_prefix))

												if i > 0:
																		# no attention layers in up_blocks.0
																		lowerCAmelCase__  :  Any						     =	F"""up_blocks.{i}.attentions.{j}."""
																		lowerCAmelCase__  :  List[str]						     =	F"""output_blocks.{3*i + j}.1."""
																		unet_conversion_map_layer.append((sd_up_atn_prefix, hf_up_atn_prefix))

						if i < 3:
												# no downsample in down_blocks.3
												lowerCAmelCase__  :  List[str]						     =	F"""down_blocks.{i}.downsamplers.0.conv."""
												lowerCAmelCase__  :  Union[str, Any]						     =	F"""input_blocks.{3*(i+1)}.0.op."""
												unet_conversion_map_layer.append((sd_downsample_prefix, hf_downsample_prefix))

												# no upsample in up_blocks.3
												lowerCAmelCase__  :  Dict						     =	F"""up_blocks.{i}.upsamplers.0."""
												lowerCAmelCase__  :  List[Any]						     =	F"""output_blocks.{3*i + 2}.{1 if i == 0 else 2}."""
												unet_conversion_map_layer.append((sd_upsample_prefix, hf_upsample_prefix))

lowerCAmelCase__  :  str						     =	'mid_block.attentions.0.'
lowerCAmelCase__  :  Union[str, Any]						     =	'middle_block.1.'
unet_conversion_map_layer.append((sd_mid_atn_prefix, hf_mid_atn_prefix))

for j in range(2):
						lowerCAmelCase__  :  int						     =	F"""mid_block.resnets.{j}."""
						lowerCAmelCase__  :  List[str]						     =	F"""middle_block.{2*j}."""
						unet_conversion_map_layer.append((sd_mid_res_prefix, hf_mid_res_prefix))



def   a_						(					lowerCamelCase       ):
				# buyer beware: this is a *brittle* function,
				# and correct output requires that all of these pieces interact in
				# the exact order in which I have arranged them.
				UpperCAmelCase__             =							{k: k for k in unet_state_dict.keys()}
				for sd_name, hf_name in unet_conversion_map:
								UpperCAmelCase__             =							sd_name
				for k, v in mapping.items():
								if "resnets" in k:
												for sd_part, hf_part in unet_conversion_map_resnet:
																UpperCAmelCase__             =							v.replace(lowerCamelCase					,			lowerCamelCase       )
												UpperCAmelCase__             =							v
				for k, v in mapping.items():
								for sd_part, hf_part in unet_conversion_map_layer:
												UpperCAmelCase__             =							v.replace(lowerCamelCase					,			lowerCamelCase       )
								UpperCAmelCase__             =							v
				UpperCAmelCase__             =							{v: unet_state_dict[k] for k, v in mapping.items()}
				return new_state_dict


# ================#
# VAE Conversion #
# ================#

lowerCAmelCase__  :  str						     =	[
    # (stable-diffusion, HF Diffusers)
    ('nin_shortcut', 'conv_shortcut'),
    ('norm_out', 'conv_norm_out'),
    ('mid.attn_1.', 'mid_block.attentions.0.'),
]

for i in range(4):
						# down_blocks have two resnets
						for j in range(2):
												lowerCAmelCase__  :  List[str]						     =	F"""encoder.down_blocks.{i}.resnets.{j}."""
												lowerCAmelCase__  :  List[Any]						     =	F"""encoder.down.{i}.block.{j}."""
												vae_conversion_map.append((sd_down_prefix, hf_down_prefix))

						if i < 3:
												lowerCAmelCase__  :  Dict						     =	F"""down_blocks.{i}.downsamplers.0."""
												lowerCAmelCase__  :  str						     =	F"""down.{i}.downsample."""
												vae_conversion_map.append((sd_downsample_prefix, hf_downsample_prefix))

												lowerCAmelCase__  :  int						     =	F"""up_blocks.{i}.upsamplers.0."""
												lowerCAmelCase__  :  str						     =	F"""up.{3-i}.upsample."""
												vae_conversion_map.append((sd_upsample_prefix, hf_upsample_prefix))

						# up_blocks have three resnets
						# also, up blocks in hf are numbered in reverse from sd
						for j in range(3):
												lowerCAmelCase__  :  Dict						     =	F"""decoder.up_blocks.{i}.resnets.{j}."""
												lowerCAmelCase__  :  Optional[int]						     =	F"""decoder.up.{3-i}.block.{j}."""
												vae_conversion_map.append((sd_up_prefix, hf_up_prefix))

# this part accounts for mid blocks in both the encoder and the decoder
for i in range(2):
						lowerCAmelCase__  :  Any						     =	F"""mid_block.resnets.{i}."""
						lowerCAmelCase__  :  Any						     =	F"""mid.block_{i+1}."""
						vae_conversion_map.append((sd_mid_res_prefix, hf_mid_res_prefix))


lowerCAmelCase__  :  str						     =	[
    # (stable-diffusion, HF Diffusers)
    ('norm.', 'group_norm.'),
    ('q.', 'query.'),
    ('k.', 'key.'),
    ('v.', 'value.'),
    ('proj_out.', 'proj_attn.'),
]



def   a_						(					lowerCamelCase       ):
				# convert HF linear weights to SD conv2d weights
				return w.reshape(*w.shape					,			1					,			1       )



def   a_						(					lowerCamelCase       ):
				UpperCAmelCase__             =							{k: k for k in vae_state_dict.keys()}
				for k, v in mapping.items():
								for sd_part, hf_part in vae_conversion_map:
												UpperCAmelCase__             =							v.replace(lowerCamelCase					,			lowerCamelCase       )
								UpperCAmelCase__             =							v
				for k, v in mapping.items():
								if "attentions" in k:
												for sd_part, hf_part in vae_conversion_map_attn:
																UpperCAmelCase__             =							v.replace(lowerCamelCase					,			lowerCamelCase       )
												UpperCAmelCase__             =							v
				UpperCAmelCase__             =							{v: vae_state_dict[k] for k, v in mapping.items()}
				UpperCAmelCase__             =							['q', 'k', 'v', 'proj_out']
				for k, v in new_state_dict.items():
								for weight_name in weights_to_convert:
												if f'''mid.attn_1.{weight_name}.weight''' in k:
																print(f'''Reshaping {k} for SD format'''       )
																UpperCAmelCase__             =							reshape_weight_for_sd(lowerCamelCase       )
				return new_state_dict


# =========================#
# Text Encoder Conversion #
# =========================#


lowerCAmelCase__  :  Optional[int]						     =	[
    # (stable-diffusion, HF Diffusers)
    ('resblocks.', 'text_model.encoder.layers.'),
    ('ln_1', 'layer_norm1'),
    ('ln_2', 'layer_norm2'),
    ('.c_fc.', '.fc1.'),
    ('.c_proj.', '.fc2.'),
    ('.attn', '.self_attn'),
    ('ln_final.', 'transformer.text_model.final_layer_norm.'),
    ('token_embedding.weight', 'transformer.text_model.embeddings.token_embedding.weight'),
    ('positional_embedding', 'transformer.text_model.embeddings.position_embedding.weight'),
]
lowerCAmelCase__  :  List[Any]						     =	{re.escape(x[1]): x[0] for x in textenc_conversion_lst}
lowerCAmelCase__  :  int						     =	re.compile('|'.join(protected.keys()))

# Ordering is from https://github.com/pytorch/pytorch/blob/master/test/cpp/api/modules.cpp
lowerCAmelCase__  :  Optional[int]						     =	{'q': 0, 'k': 1, 'v': 2}



def   a_						(					lowerCamelCase       ):
				UpperCAmelCase__             =							{}
				UpperCAmelCase__             =							{}
				UpperCAmelCase__             =							{}
				for k, v in text_enc_dict.items():
								if (
								    k.endswith('.self_attn.q_proj.weight'       )
								    or k.endswith('.self_attn.k_proj.weight'       )
								    or k.endswith('.self_attn.v_proj.weight'       )
								):
												UpperCAmelCase__             =							k[: -len('.q_proj.weight'       )]
												UpperCAmelCase__             =							k[-len('q_proj.weight'       )]
												if k_pre not in capture_qkv_weight:
																UpperCAmelCase__             =							[None, None, None]
												UpperCAmelCase__             =							v
												continue

								if (
								    k.endswith('.self_attn.q_proj.bias'       )
								    or k.endswith('.self_attn.k_proj.bias'       )
								    or k.endswith('.self_attn.v_proj.bias'       )
								):
												UpperCAmelCase__             =							k[: -len('.q_proj.bias'       )]
												UpperCAmelCase__             =							k[-len('q_proj.bias'       )]
												if k_pre not in capture_qkv_bias:
																UpperCAmelCase__             =							[None, None, None]
												UpperCAmelCase__             =							v
												continue

								UpperCAmelCase__             =							textenc_pattern.sub(lambda lowerCamelCase       : protected[re.escape(m.group(0       )       )]					,			lowerCamelCase       )
								UpperCAmelCase__             =							v

				for k_pre, tensors in capture_qkv_weight.items():
								if None in tensors:
												raise Exception('CORRUPTED MODEL: one of the q-k-v values for the text encoder was missing'       )
								UpperCAmelCase__             =							textenc_pattern.sub(lambda lowerCamelCase       : protected[re.escape(m.group(0       )       )]					,			lowerCamelCase       )
								UpperCAmelCase__             =							torch.cat(lowerCamelCase       )

				for k_pre, tensors in capture_qkv_bias.items():
								if None in tensors:
												raise Exception('CORRUPTED MODEL: one of the q-k-v values for the text encoder was missing'       )
								UpperCAmelCase__             =							textenc_pattern.sub(lambda lowerCamelCase       : protected[re.escape(m.group(0       )       )]					,			lowerCamelCase       )
								UpperCAmelCase__             =							torch.cat(lowerCamelCase       )

				return new_state_dict



def   a_						(					lowerCamelCase       ):
				return text_enc_dict


if __name__ == "__main__":
						lowerCAmelCase__  :  Tuple						     =	argparse.ArgumentParser()

						parser.add_argument('--model_path', default=None, type=str, required=True, help='Path to the model to convert.')
						parser.add_argument('--checkpoint_path', default=None, type=str, required=True, help='Path to the output model.')
						parser.add_argument('--half', action='store_true', help='Save weights in half precision.')
						parser.add_argument(
						    '--use_safetensors', action='store_true', help='Save weights use safetensors, default is ckpt.'
						)

						lowerCAmelCase__  :  Optional[int]						     =	parser.parse_args()

						assert args.model_path is not None, "Must provide a model path!"

						assert args.checkpoint_path is not None, "Must provide a checkpoint path!"

						# Path for safetensors
						lowerCAmelCase__  :  Tuple						     =	osp.join(args.model_path, 'unet', 'diffusion_pytorch_model.safetensors')
						lowerCAmelCase__  :  List[str]						     =	osp.join(args.model_path, 'vae', 'diffusion_pytorch_model.safetensors')
						lowerCAmelCase__  :  int						     =	osp.join(args.model_path, 'text_encoder', 'model.safetensors')

						# Load models from safetensors if it exists, if it doesn't pytorch
						if osp.exists(unet_path):
												lowerCAmelCase__  :  Union[str, Any]						     =	load_file(unet_path, device='cpu')
						else:
												lowerCAmelCase__  :  str						     =	osp.join(args.model_path, 'unet', 'diffusion_pytorch_model.bin')
												lowerCAmelCase__  :  Dict						     =	torch.load(unet_path, map_location='cpu')

						if osp.exists(vae_path):
												lowerCAmelCase__  :  Optional[Any]						     =	load_file(vae_path, device='cpu')
						else:
												lowerCAmelCase__  :  Optional[int]						     =	osp.join(args.model_path, 'vae', 'diffusion_pytorch_model.bin')
												lowerCAmelCase__  :  List[str]						     =	torch.load(vae_path, map_location='cpu')

						if osp.exists(text_enc_path):
												lowerCAmelCase__  :  Tuple						     =	load_file(text_enc_path, device='cpu')
						else:
												lowerCAmelCase__  :  Any						     =	osp.join(args.model_path, 'text_encoder', 'pytorch_model.bin')
												lowerCAmelCase__  :  Any						     =	torch.load(text_enc_path, map_location='cpu')

						# Convert the UNet model
						lowerCAmelCase__  :  Any						     =	convert_unet_state_dict(unet_state_dict)
						lowerCAmelCase__  :  Dict						     =	{'model.diffusion_model.' + k: v for k, v in unet_state_dict.items()}

						# Convert the VAE model
						lowerCAmelCase__  :  List[Any]						     =	convert_vae_state_dict(vae_state_dict)
						lowerCAmelCase__  :  str						     =	{'first_stage_model.' + k: v for k, v in vae_state_dict.items()}

						# Easiest way to identify v2.0 model seems to be that the text encoder (OpenCLIP) is deeper
						lowerCAmelCase__  :  List[Any]						     =	'text_model.encoder.layers.22.layer_norm2.bias' in text_enc_dict

						if is_vaa_model:
												# Need to add the tag 'transformer' in advance so we can knock it out from the final layer-norm
												lowerCAmelCase__  :  Tuple						     =	{'transformer.' + k: v for k, v in text_enc_dict.items()}
												lowerCAmelCase__  :  List[str]						     =	convert_text_enc_state_dict_vaa(text_enc_dict)
												lowerCAmelCase__  :  str						     =	{'cond_stage_model.model.' + k: v for k, v in text_enc_dict.items()}
						else:
												lowerCAmelCase__  :  Optional[Any]						     =	convert_text_enc_state_dict(text_enc_dict)
												lowerCAmelCase__  :  Optional[Any]						     =	{'cond_stage_model.transformer.' + k: v for k, v in text_enc_dict.items()}

						# Put together new checkpoint
						lowerCAmelCase__  :  List[Any]						     =	{**unet_state_dict, **vae_state_dict, **text_enc_dict}
						if args.half:
												lowerCAmelCase__  :  int						     =	{k: v.half() for k, v in state_dict.items()}

						if args.use_safetensors:
												save_file(state_dict, args.checkpoint_path)
						else:
												lowerCAmelCase__  :  List[str]						     =	{'state_dict': state_dict}
												torch.save(state_dict, args.checkpoint_path)



