'''simple docstring'''

import argparse
import os.path as osp
import re

import torch
from safetensors.torch import load_file, save_file


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

lowerCAmelCase_ : 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_ : List[str]						  =      []
# 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_ : List[str]						  =      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_ : Union[str, Any]						  =      F'''down_blocks.{i}.attentions.{j}.'''
                     lowerCAmelCase_ : List[Any]						  =      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_ : Any						  =      F'''up_blocks.{i}.resnets.{j}.'''
              lowerCAmelCase_ : Optional[Any]						  =      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_ : Union[str, Any]						  =      F'''up_blocks.{i}.attentions.{j}.'''
                     lowerCAmelCase_ : Union[str, Any]						  =      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_ : Dict						  =      F'''down_blocks.{i}.downsamplers.0.conv.'''
              lowerCAmelCase_ : int						  =      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_ : Optional[int]						  =      F'''up_blocks.{i}.upsamplers.0.'''
              lowerCAmelCase_ : Dict						  =      F'''output_blocks.{3*i + 2}.{1 if i == 0 else 2}.'''
              unet_conversion_map_layer.append((sd_upsample_prefix, hf_upsample_prefix))

lowerCAmelCase_ : Optional[int]						  =      "mid_block.attentions.0."
lowerCAmelCase_ : str						  =      "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			:						List[str]				)	->							List[Any]:
      '''simple docstring'''


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


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

lowerCAmelCase_ : Optional[Any]						  =      [
    # (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_ : Optional[Any]						  =      F'''encoder.down_blocks.{i}.resnets.{j}.'''
              lowerCAmelCase_ : Optional[int]						  =      F'''encoder.down.{i}.block.{j}.'''
              vae_conversion_map.append((sd_down_prefix, hf_down_prefix))

       if i < 3:
              lowerCAmelCase_ : int						  =      F'''down_blocks.{i}.downsamplers.0.'''
              lowerCAmelCase_ : str						  =      F'''down.{i}.downsample.'''
              vae_conversion_map.append((sd_downsample_prefix, hf_downsample_prefix))

              lowerCAmelCase_ : Dict						  =      F'''up_blocks.{i}.upsamplers.0.'''
              lowerCAmelCase_ : int						  =      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_ : str						  =      F'''decoder.up_blocks.{i}.resnets.{j}.'''
              lowerCAmelCase_ : Dict						  =      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_ : Tuple						  =      F'''mid_block.resnets.{i}.'''
       lowerCAmelCase_ : Optional[Any]						  =      F'''mid.block_{i+1}.'''
       vae_conversion_map.append((sd_mid_res_prefix, hf_mid_res_prefix))


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


def __a (	__lowerCamelCase			:						List[Any]				)	->							Tuple:
      '''simple docstring'''


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


def __a (	__lowerCamelCase			:						Union[str, Any]				)	->							str:
      '''simple docstring'''


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


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


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

# Ordering is from https://github.com/pytorch/pytorch/blob/master/test/cpp/api/modules.cpp
lowerCAmelCase_ : Union[str, Any]						  =      {"q": 0, "k": 1, "v": 2}


def __a (	__lowerCamelCase			:						int				)	->							List[str]:
      '''simple docstring'''


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

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

      return new_state_dict


def __a (	__lowerCamelCase			:						Optional[Any]				)	->							int:
      '''simple docstring'''


      return text_enc_dict


if __name__ == "__main__":
       lowerCAmelCase_ : int						  =      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[Any]						  =      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_ : Any						  =      osp.join(args.model_path, "unet", "diffusion_pytorch_model.safetensors")
       lowerCAmelCase_ : Any						  =      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_ : Optional[Any]						  =      load_file(unet_path, device="cpu")
       else:
              lowerCAmelCase_ : Optional[Any]						  =      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[int]						  =      load_file(vae_path, device="cpu")
       else:
              lowerCAmelCase_ : Any						  =      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_ : Dict						  =      load_file(text_enc_path, device="cpu")
       else:
              lowerCAmelCase_ : Optional[Any]						  =      osp.join(args.model_path, "text_encoder", "pytorch_model.bin")
              lowerCAmelCase_ : int						  =      torch.load(text_enc_path, map_location="cpu")

       # Convert the UNet model
       lowerCAmelCase_ : Optional[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_ : Optional[Any]						  =      convert_vae_state_dict(vae_state_dict)
       lowerCAmelCase_ : Optional[int]						  =      {"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_ : 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_ : List[Any]						  =      {"transformer." + k: v for k, v in text_enc_dict.items()}
              lowerCAmelCase_ : int						  =      convert_text_enc_state_dict_vaa(text_enc_dict)
              lowerCAmelCase_ : Any						  =      {"cond_stage_model.model." + k: v for k, v in text_enc_dict.items()}
       else:
              lowerCAmelCase_ : List[str]						  =      convert_text_enc_state_dict(text_enc_dict)
              lowerCAmelCase_ : Dict						  =      {"cond_stage_model.transformer." + k: v for k, v in text_enc_dict.items()}

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

       if args.use_safetensors:
              save_file(state_dict, args.checkpoint_path)
       else:
              lowerCAmelCase_ : int						  =      {"state_dict": state_dict}
              torch.save(state_dict, args.checkpoint_path)







'''simple docstring'''

import io
import itertools
import json
from dataclasses import dataclass
from typing import Optional

import pyarrow as pa
import pyarrow.json as paj

import datasets
from datasets.table import table_cast
from datasets.utils.file_utils import readline


lowerCAmelCase_ : Union[str, Any]						  =      datasets.utils.logging.get_logger(__name__)





@dataclass
class 							lowercase						(						datasets.BuilderConfig ):
   lowerCamelCase_    					=None
   lowerCamelCase_    					="utf-8"
   lowerCamelCase_    					=None
   lowerCamelCase_    					=None
   lowerCamelCase_    					=True  # deprecated
   lowerCamelCase_    					=None  # deprecated
   lowerCamelCase_    					=1_0 << 2_0  # 10MB
   lowerCamelCase_    					=None





class 							lowercase						(						datasets.ArrowBasedBuilder ):
   lowerCamelCase_    					=JsonConfig


   def __UpperCAmelCase       (   self				:	int)      ->					List[str]:
         if self.config.block_size is not None:
               logger.warning("The JSON loader parameter `block_size` is deprecated. Please use `chunksize` instead")
               lowercase_     =      self.config.block_size
         if self.config.use_threads is not True:
               logger.warning(
                   "The JSON loader parameter `use_threads` is deprecated and doesn't have any effect anymore.")
         if self.config.newlines_in_values is not None:
               raise ValueError("The JSON loader parameter `newlines_in_values` is no longer supported")
         return datasets.DatasetInfo(features=self.config.features)


   def __UpperCAmelCase       (   self				:	Any   ,		__lowerCAmelCase				:	int)      ->					Any:
         if not self.config.data_files:
               raise ValueError(F'At least one data file must be specified, but got data_files={self.config.data_files}')
         lowercase_     =      dl_manager.download_and_extract(self.config.data_files)
         if isinstance(__lowerCAmelCase   ,		(str, list, tuple)):
               lowercase_     =      data_files
               if isinstance(__lowerCAmelCase   ,		__lowerCAmelCase):
                     lowercase_     =      [files]
               lowercase_     =      [dl_manager.iter_files(__lowerCAmelCase) for file in files]
               return [datasets.SplitGenerator(name=datasets.Split.TRAIN   ,		gen_kwargs={"files": files})]
         lowercase_     =      []
         for split_name, files in data_files.items():
               if isinstance(__lowerCAmelCase   ,		__lowerCAmelCase):
                     lowercase_     =      [files]
               lowercase_     =      [dl_manager.iter_files(__lowerCAmelCase) for file in files]
               splits.append(datasets.SplitGenerator(name=__lowerCAmelCase   ,		gen_kwargs={"files": files}))
         return splits


   def __UpperCAmelCase       (   self				:	Any   ,		__lowerCAmelCase				:	pa.Table)      ->					pa.Table:
         if self.config.features is not None:
               # adding missing columns
               for column_name in set(self.config.features) - set(pa_table.column_names):
                     lowercase_     =      self.config.features.arrow_schema.field(__lowerCAmelCase).type
                     lowercase_     =      pa_table.append_column(__lowerCAmelCase   ,		pa.array([None] * len(__lowerCAmelCase)   ,		type=__lowerCAmelCase))
               # more expensive cast to support nested structures with keys in a different order
               # allows str <-> int/float or str to Audio for example
               lowercase_     =      table_cast(__lowerCAmelCase   ,		self.config.features.arrow_schema)
         return pa_table


   def __UpperCAmelCase       (   self				:	List[Any]   ,		__lowerCAmelCase				:	List[Any])      ->					List[Any]:
         for file_idx, file in enumerate(itertools.chain.from_iterable(__lowerCAmelCase)):
               # If the file is one json object and if we need to look at the list of items in one specific field
               if self.config.field is not None:
                     with open(__lowerCAmelCase   ,		encoding=self.config.encoding   ,		errors=self.config.encoding_errors) as f:
                           lowercase_     =      json.load(__lowerCAmelCase)

                     # We keep only the field we are interested in
                     lowercase_     =      dataset[self.config.field]

                     # We accept two format: a list of dicts or a dict of lists
                     if isinstance(__lowerCAmelCase   ,		(list, tuple)):
                           lowercase_     =      set().union(*[row.keys() for row in dataset])
                           lowercase_     =      {col: [row.get(__lowerCAmelCase) for row in dataset] for col in keys}
                     else:
                           lowercase_     =      dataset
                     lowercase_     =      pa.Table.from_pydict(__lowerCAmelCase)
                     yield file_idx, self._cast_table(__lowerCAmelCase)

               # If the file has one json object per line
               else:
                     with open(__lowerCAmelCase   ,		"rb") as f:
                           lowercase_     =      0
                           # Use block_size equal to the chunk size divided by 32 to leverage multithreading
                           # Set a default minimum value of 16kB if the chunk size is really small
                           lowercase_     =      max(self.config.chunksize // 32   ,		16 << 10)
                           lowercase_     =      (
                               self.config.encoding_errors if self.config.encoding_errors is not None else "strict"
                           )
                           while True:
                                 lowercase_     =      f.read(self.config.chunksize)
                                 if not batch:
                                       break
                                 # Finish current line
                                 try:
                                       batch += f.readline()
                                 except (AttributeError, io.UnsupportedOperation):
                                       batch += readline(__lowerCAmelCase)
                                 # PyArrow only accepts utf-8 encoded bytes
                                 if self.config.encoding != "utf-8":
                                       lowercase_     =      batch.decode(self.config.encoding   ,		errors=__lowerCAmelCase).encode("utf-8")
                                 try:
                                       while True:
                                             try:
                                                   lowercase_     =      paj.read_json(
                                                       io.BytesIO(__lowerCAmelCase)   ,		read_options=paj.ReadOptions(block_size=__lowerCAmelCase))
                                                   break
                                             except (pa.ArrowInvalid, pa.ArrowNotImplementedError) as e:
                                                   if (
                                                       isinstance(__lowerCAmelCase   ,		pa.ArrowInvalid)
                                                       and "straddling" not in str(__lowerCAmelCase)
                                                       or block_size > len(__lowerCAmelCase)
                                                   ):
                                                         raise
                                                   else:
                                                         # Increase the block size in case it was too small.
                                                         # The block size will be reset for the next file.
                                                         logger.debug(
                                                             F'Batch of {len(__lowerCAmelCase)} bytes couldn\'t be parsed with block_size={block_size}. Retrying with block_size={block_size * 2}.')
                                                         block_size *= 2
                                 except pa.ArrowInvalid as e:
                                       try:
                                             with open(
                                                 __lowerCAmelCase   ,		encoding=self.config.encoding   ,		errors=self.config.encoding_errors) as f:
                                                   lowercase_     =      json.load(__lowerCAmelCase)
                                       except json.JSONDecodeError:
                                             logger.error(F'Failed to read file \'{file}\' with error {type(__lowerCAmelCase)}: {e}')
                                             raise e
                                       # If possible, parse the file as a list of json objects and exit the loop
                                       if isinstance(__lowerCAmelCase   ,		__lowerCAmelCase):  # list is the only sequence type supported in JSON
                                             try:
                                                   lowercase_     =      set().union(*[row.keys() for row in dataset])
                                                   lowercase_     =      {col: [row.get(__lowerCAmelCase) for row in dataset] for col in keys}
                                                   lowercase_     =      pa.Table.from_pydict(__lowerCAmelCase)
                                             except (pa.ArrowInvalid, AttributeError) as e:
                                                   logger.error(F'Failed to read file \'{file}\' with error {type(__lowerCAmelCase)}: {e}')
                                                   raise ValueError(F'Not able to read records in the JSON file at {file}.') from None
                                             yield file_idx, self._cast_table(__lowerCAmelCase)
                                             break
                                       else:
                                             logger.error(F'Failed to read file \'{file}\' with error {type(__lowerCAmelCase)}: {e}')
                                             raise ValueError(
                                                 F'Not able to read records in the JSON file at {file}. '
                                                 F'You should probably indicate the field of the JSON file containing your records. '
                                                 F'This JSON file contain the following fields: {str(list(dataset.keys()))}. '
                                                 F'Select the correct one and provide it as `field=\'XXX\'` to the dataset loading method. ') from None
                        # Uncomment for debugging (will print the Arrow table size and elements)
                        # logger.warning(f"pa_table: {pa_table} num rows: {pa_table.num_rows}")
                        # logger.warning('\n'.join(str(pa_table.slice(i, 1).to_pydict()) for i in range(pa_table.num_rows)))
                                 yield (file_idx, batch_idx), self._cast_table(__lowerCAmelCase)
                                 batch_idx += 1






