






"""simple docstring"""




import argparse
import os.path as osp
import re

import torch
from safetensors.torch import load_file, save_file


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

UpperCamelCase							       =  [
    # (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'),
]

UpperCamelCase							       =  [
    # (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'),
]

UpperCamelCase							       =  []
# 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
          UpperCamelCase							       =  f"down_blocks.{i}.resnets.{j}."
          UpperCamelCase							       =  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
               UpperCamelCase							       =  f"down_blocks.{i}.attentions.{j}."
               UpperCamelCase							       =  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
          UpperCamelCase							       =  f"up_blocks.{i}.resnets.{j}."
          UpperCamelCase							       =  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
               UpperCamelCase							       =  f"up_blocks.{i}.attentions.{j}."
               UpperCamelCase							       =  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
          UpperCamelCase							       =  f"down_blocks.{i}.downsamplers.0.conv."
          UpperCamelCase							       =  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
          UpperCamelCase							       =  f"up_blocks.{i}.upsamplers.0."
          UpperCamelCase							       =  f"output_blocks.{3*i + 2}.{1 if i == 0 else 2}."
          unet_conversion_map_layer.append((sd_upsample_prefix, hf_upsample_prefix))

UpperCamelCase							       =  'mid_block.attentions.0.'
UpperCamelCase							       =  'middle_block.1.'
unet_conversion_map_layer.append((sd_mid_atn_prefix, hf_mid_atn_prefix))

for j in range(2):
     UpperCamelCase							       =  f"mid_block.resnets.{j}."
     UpperCamelCase							       =  f"middle_block.{2*j}."
     unet_conversion_map_layer.append((sd_mid_res_prefix, hf_mid_res_prefix))






def        lowerCAmelCase (      UpperCamelCase_:							int					)       ->							List[Any]:







  '''simple docstring'''







  _a							=	{k: k for k in unet_state_dict.keys()}
  for sd_name, hf_name in unet_conversion_map:
    _a							=	sd_name
  for k, v in mapping.items():
    if "resnets" in k:
      for sd_part, hf_part in unet_conversion_map_resnet:
        _a							=	v.replace(snake_case_     , snake_case_					)
      _a							=	v
  for k, v in mapping.items():
    for sd_part, hf_part in unet_conversion_map_layer:
      _a							=	v.replace(snake_case_     , snake_case_					)
    _a							=	v
  _a							=	{v: unet_state_dict[k] for k, v in mapping.items()}
  return new_state_dict


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

UpperCamelCase							       =  [
    # (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):
          UpperCamelCase							       =  f"encoder.down_blocks.{i}.resnets.{j}."
          UpperCamelCase							       =  f"encoder.down.{i}.block.{j}."
          vae_conversion_map.append((sd_down_prefix, hf_down_prefix))

     if i < 3:
          UpperCamelCase							       =  f"down_blocks.{i}.downsamplers.0."
          UpperCamelCase							       =  f"down.{i}.downsample."
          vae_conversion_map.append((sd_downsample_prefix, hf_downsample_prefix))

          UpperCamelCase							       =  f"up_blocks.{i}.upsamplers.0."
          UpperCamelCase							       =  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):
          UpperCamelCase							       =  f"decoder.up_blocks.{i}.resnets.{j}."
          UpperCamelCase							       =  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):
     UpperCamelCase							       =  f"mid_block.resnets.{i}."
     UpperCamelCase							       =  f"mid.block_{i+1}."
     vae_conversion_map.append((sd_mid_res_prefix, hf_mid_res_prefix))


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






def        lowerCAmelCase (      UpperCamelCase_:							Optional[Any]					)       ->							Union[str, Any]:







  '''simple docstring'''







  return w.reshape(*w.shape     , 1     , 1					)






def        lowerCAmelCase (      UpperCamelCase_:							Tuple					)       ->							Union[str, Any]:







  '''simple docstring'''







  _a							=	{k: k for k in vae_state_dict.keys()}
  for k, v in mapping.items():
    for sd_part, hf_part in vae_conversion_map:
      _a							=	v.replace(snake_case_     , snake_case_					)
    _a							=	v
  for k, v in mapping.items():
    if "attentions" in k:
      for sd_part, hf_part in vae_conversion_map_attn:
        _a							=	v.replace(snake_case_     , snake_case_					)
      _a							=	v
  _a							=	{v: vae_state_dict[k] for k, v in mapping.items()}
  _a							=	['''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'''					)
        _a							=	reshape_weight_for_sd(snake_case_					)
  return new_state_dict


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


UpperCamelCase							       =  [
    # (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'),
]
UpperCamelCase							       =  {re.escape(x[1]): x[0] for x in textenc_conversion_lst}
UpperCamelCase							       =  re.compile("""|""".join(protected.keys()))

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






def        lowerCAmelCase (      UpperCamelCase_:							List[str]					)       ->							List[str]:







  '''simple docstring'''







  _a							=	{}
  _a							=	{}
  _a							=	{}
  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"					)
    ):
      _a							=	k[: -len(".q_proj.weight"					)]
      _a							=	k[-len("q_proj.weight"					)]
      if k_pre not in capture_qkv_weight:
        _a							=	[None, None, None]
      _a							=	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"					)
    ):
      _a							=	k[: -len(".q_proj.bias"					)]
      _a							=	k[-len("q_proj.bias"					)]
      if k_pre not in capture_qkv_bias:
        _a							=	[None, None, None]
      _a							=	v
      continue

    _a							=	textenc_pattern.sub(lambda UpperCamelCase_					: protected[re.escape(m.group(0					)					)]     , snake_case_					)
    _a							=	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"					)
    _a							=	textenc_pattern.sub(lambda UpperCamelCase_					: protected[re.escape(m.group(0					)					)]     , snake_case_					)
    _a							=	torch.cat(snake_case_					)

  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"					)
    _a							=	textenc_pattern.sub(lambda UpperCamelCase_					: protected[re.escape(m.group(0					)					)]     , snake_case_					)
    _a							=	torch.cat(snake_case_					)

  return new_state_dict






def        lowerCAmelCase (      UpperCamelCase_:							str					)       ->							List[Any]:







  '''simple docstring'''







  return text_enc_dict


if __name__ == "__main__":
     UpperCamelCase							       =  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."""
     )

     UpperCamelCase							       =  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
     UpperCamelCase							       =  osp.join(args.model_path, """unet""", """diffusion_pytorch_model.safetensors""")
     UpperCamelCase							       =  osp.join(args.model_path, """vae""", """diffusion_pytorch_model.safetensors""")
     UpperCamelCase							       =  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):
          UpperCamelCase							       =  load_file(unet_path, device="""cpu""")
     else:
          UpperCamelCase							       =  osp.join(args.model_path, """unet""", """diffusion_pytorch_model.bin""")
          UpperCamelCase							       =  torch.load(unet_path, map_location="""cpu""")

     if osp.exists(vae_path):
          UpperCamelCase							       =  load_file(vae_path, device="""cpu""")
     else:
          UpperCamelCase							       =  osp.join(args.model_path, """vae""", """diffusion_pytorch_model.bin""")
          UpperCamelCase							       =  torch.load(vae_path, map_location="""cpu""")

     if osp.exists(text_enc_path):
          UpperCamelCase							       =  load_file(text_enc_path, device="""cpu""")
     else:
          UpperCamelCase							       =  osp.join(args.model_path, """text_encoder""", """pytorch_model.bin""")
          UpperCamelCase							       =  torch.load(text_enc_path, map_location="""cpu""")

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

     # Convert the VAE model
     UpperCamelCase							       =  convert_vae_state_dict(vae_state_dict)
     UpperCamelCase							       =  {'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
     UpperCamelCase							       =  '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
          UpperCamelCase							       =  {'transformer.' + k: v for k, v in text_enc_dict.items()}
          UpperCamelCase							       =  convert_text_enc_state_dict_vaa(text_enc_dict)
          UpperCamelCase							       =  {'cond_stage_model.model.' + k: v for k, v in text_enc_dict.items()}
     else:
          UpperCamelCase							       =  convert_text_enc_state_dict(text_enc_dict)
          UpperCamelCase							       =  {'cond_stage_model.transformer.' + k: v for k, v in text_enc_dict.items()}

     # Put together new checkpoint
     UpperCamelCase							       =  {**unet_state_dict, **vae_state_dict, **text_enc_dict}
     if args.half:
          UpperCamelCase							       =  {k: v.half() for k, v in state_dict.items()}

     if args.use_safetensors:
          save_file(state_dict, args.checkpoint_path)
     else:
          UpperCamelCase							       =  {'state_dict': state_dict}
          torch.save(state_dict, args.checkpoint_path)










"""simple docstring"""




import numpy as np
from numpy import ndarray
from scipy.optimize import Bounds, LinearConstraint, minimize






def        lowerCAmelCase (      UpperCamelCase_:							ndarray					)       ->							float:







  '''simple docstring'''







  return np.dot(UpperCamelCase_     , UpperCamelCase_					)




class    lowercase_    :






      def __init__(    self     ,     *,
          a_ = np.inf     ,     a_ = "linear"     ,     a_ = 0.0     ,     )	->None:







        '''simple docstring'''


        _a							=	regularization
        _a							=	gamma
        if kernel == "linear":
          _a							=	self.__linear
        elif kernel == "rbf":
          if self.gamma == 0:
            raise ValueError("rbf kernel requires gamma"	)
          if not isinstance(self.gamma     ,     (float, int)	):
            raise ValueError("gamma must be float or int"	)
          if not self.gamma > 0:
            raise ValueError("gamma must be > 0"	)
          _a							=	self.__rbf
          # in the future, there could be a default value like in sklearn
          # sklear: def_gamma = 1/(n_features * X.var()) (wiki)
          # previously it was 1/(n_features)
        else:
          _a							=	f'''Unknown kernel: {kernel}'''
          raise ValueError(a_	)






      def 			lowerCamelCase__    (    self     ,     a_     ,     a_	)	->float:







        '''simple docstring'''


        return np.dot(a_     ,     a_	)






      def 			lowerCamelCase__    (    self     ,     a_     ,     a_	)	->float:







        '''simple docstring'''


        return np.exp(-(self.gamma * norm_squared(vectora - vectora	))	)






      def 			lowerCamelCase__    (    self     ,     a_     ,     a_	)	->None:







        '''simple docstring'''


        _a							=	observations
        _a							=	classes

        # using Wolfe's Dual to calculate w.
        # Primal problem: minimize 1/2*norm_squared(w)
        #   constraint: yn(w . xn + b) >= 1
        #
        # With l a vector
        # Dual problem: maximize sum_n(ln) -
        #       1/2 * sum_n(sum_m(ln*lm*yn*ym*xn . xm))
        #   constraint: self.C >= ln >= 0
        #           and sum_n(ln*yn) = 0
        # Then we get w using w = sum_n(ln*yn*xn)
        # At the end we can get b ~= mean(yn - w . xn)
        #
        # Since we use kernels, we only need l_star to calculate b
        # and to classify observations

        ((_a)			,	)							=	np.shape(a_	)

        def to_minimize(a_	) -> float:
          _a							=	0
          ((_a)			,	)							=	np.shape(a_	)
          for i in range(a_	):
            for j in range(a_	):
              s += (
                  candidate[i]
                  * candidate[j]
                  * classes[i]
                  * classes[j]
                  * self.kernel(observations[i]     ,     observations[j]	)
              )
          return 1 / 2 * s - sum(a_	)

        _a							=	LinearConstraint(a_     ,     0     ,     0	)
        _a							=	Bounds(0     ,     self.regularization	)

        _a							=	minimize(
            a_     ,     np.ones(a_	)     ,     bounds=a_     ,     constraints=[ly_contraint]	).x
        _a							=	l_star

        # calculating mean offset of separation plane to points
        _a							=	0
        for i in range(a_	):
          for j in range(a_	):
            s += classes[i] - classes[i] * self.optimum[i] * self.kernel(
                observations[i]     ,     observations[j]	)
        _a							=	s / n






      def 			lowerCamelCase__    (    self     ,     a_	)	->int:







        '''simple docstring'''


        _a							=	sum(
            self.optimum[n]
            * self.classes[n]
            * self.kernel(self.observations[n]     ,     a_	)
            for n in range(len(self.classes	)	)	)
        return 1 if s + self.offset >= 0 else -1


if __name__ == "__main__":
     import doctest

     doctest.testmod()


