


import argparse
import os.path as osp
import re

import torch
from safetensors.torch import load_file, save_file


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

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

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

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

snake_case__				:	Optional[int]					       =					'''mid_block.attentions.0.'''
snake_case__				:	List[Any]					       =					'''middle_block.1.'''
unet_conversion_map_layer.append((sd_mid_atn_prefix, hf_mid_atn_prefix))

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


def    __lowerCamelCase   (      A__						: List[str]					)     ->    int:
   # 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.
   lowerCamelCase_						:   Tuple		= {k: k for k in unet_state_dict.keys()}
   for sd_name, hf_name in unet_conversion_map:
      lowerCamelCase_						:   Optional[int]		= sd_name
   for k, v in mapping.items():
      if "resnets" in k:
         for sd_part, hf_part in unet_conversion_map_resnet:
            lowerCamelCase_						:   str		= v.replace(__snake_case						,						__snake_case					)
         lowerCamelCase_						:   Optional[int]		= v
   for k, v in mapping.items():
      for sd_part, hf_part in unet_conversion_map_layer:
         lowerCamelCase_						:   Optional[Any]		= v.replace(__snake_case						,						__snake_case					)
      lowerCamelCase_						:   Optional[int]		= v
   lowerCamelCase_						:   Dict		= {v: unet_state_dict[k] for k, v in mapping.items()}
   return new_state_dict


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

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

      if i < 3:
            snake_case__				:	int					       =					F'down_blocks.{i}.downsamplers.0.'
            snake_case__				:	List[Any]					       =					F'down.{i}.downsample.'
            vae_conversion_map.append((sd_downsample_prefix, hf_downsample_prefix))

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


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


def    __lowerCamelCase   (      A__						: List[Any]					)     ->    List[str]:
   # convert HF linear weights to SD conv2d weights
   return w.reshape(*w.shape						,						1						,						1					)


def    __lowerCamelCase   (      A__						: Dict					)     ->    Union[str, Any]:
   lowerCamelCase_						:   Tuple		= {k: k for k in vae_state_dict.keys()}
   for k, v in mapping.items():
      for sd_part, hf_part in vae_conversion_map:
         lowerCamelCase_						:   List[Any]		= v.replace(__snake_case						,						__snake_case					)
      lowerCamelCase_						:   int		= v
   for k, v in mapping.items():
      if "attentions" in k:
         for sd_part, hf_part in vae_conversion_map_attn:
            lowerCamelCase_						:   Any		= v.replace(__snake_case						,						__snake_case					)
         lowerCamelCase_						:   Optional[int]		= v
   lowerCamelCase_						:   int		= {v: vae_state_dict[k] for k, v in mapping.items()}
   lowerCamelCase_						:   Optional[int]		= ["""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'''					)
            lowerCamelCase_						:   Optional[int]		= reshape_weight_for_sd(__snake_case					)
   return new_state_dict


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


snake_case__				:	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'''),
]
snake_case__				:	List[str]					       =					{re.escape(x[1]): x[0] for x in textenc_conversion_lst}
snake_case__				:	Optional[Any]					       =					re.compile('|'.join(protected.keys()))

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


def    __lowerCamelCase   (      A__						: List[Any]					)     ->    Dict:
   lowerCamelCase_						:   Tuple		= {}
   lowerCamelCase_						:   Union[str, Any]		= {}
   lowerCamelCase_						:   str		= {}
   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"""					)
      ):
         lowerCamelCase_						:   Union[str, Any]		= k[: -len(""".q_proj.weight"""					)]
         lowerCamelCase_						:   str		= k[-len("""q_proj.weight"""					)]
         if k_pre not in capture_qkv_weight:
            lowerCamelCase_						:   Union[str, Any]		= [None, None, None]
         lowerCamelCase_						:   List[Any]		= 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"""					)
      ):
         lowerCamelCase_						:   List[str]		= k[: -len(""".q_proj.bias"""					)]
         lowerCamelCase_						:   List[Any]		= k[-len("""q_proj.bias"""					)]
         if k_pre not in capture_qkv_bias:
            lowerCamelCase_						:   List[str]		= [None, None, None]
         lowerCamelCase_						:   Dict		= v
         continue

      lowerCamelCase_						:   Optional[int]		= textenc_pattern.sub(lambda A__					: protected[re.escape(m.group(0					)					)]						,						__snake_case					)
      lowerCamelCase_						:   Union[str, Any]		= 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"""					)
      lowerCamelCase_						:   List[str]		= textenc_pattern.sub(lambda A__					: protected[re.escape(m.group(0					)					)]						,						__snake_case					)
      lowerCamelCase_						:   Optional[Any]		= 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"""					)
      lowerCamelCase_						:   Tuple		= textenc_pattern.sub(lambda A__					: protected[re.escape(m.group(0					)					)]						,						__snake_case					)
      lowerCamelCase_						:   Dict		= torch.cat(__snake_case					)

   return new_state_dict


def    __lowerCamelCase   (      A__						: Dict					)     ->    Union[str, Any]:
   return text_enc_dict


if __name__ == "__main__":
      snake_case__				:	List[str]					       =					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.'
      )

      snake_case__				:	List[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
      snake_case__				:	Any					       =					osp.join(args.model_path, 'unet', 'diffusion_pytorch_model.safetensors')
      snake_case__				:	Optional[Any]					       =					osp.join(args.model_path, 'vae', 'diffusion_pytorch_model.safetensors')
      snake_case__				:	Tuple					       =					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):
            snake_case__				:	Any					       =					load_file(unet_path, device='cpu')
      else:
            snake_case__				:	Any					       =					osp.join(args.model_path, 'unet', 'diffusion_pytorch_model.bin')
            snake_case__				:	str					       =					torch.load(unet_path, map_location='cpu')

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

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

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

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

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

      if args.use_safetensors:
            save_file(state_dict, args.checkpoint_path)
      else:
            snake_case__				:	Optional[Any]					       =					{'''state_dict''': state_dict}
            torch.save(state_dict, args.checkpoint_path)






'''simple docstring'''






import logging
import random

import ray

from transformers import RagConfig, RagRetriever, RagTokenizer
from transformers.models.rag.retrieval_rag import CustomHFIndex


_UpperCAmelCase    :     int         =       logging.getLogger(__name__)







class lowercase_       :







   """simple docstring"""

   def __init__(      self      :			Tuple     )  ->			int:
        _A    		=  False
   def 	__UpperCAmelCase   (      self      :			Optional[int], UpperCamelCase__      :			str, UpperCamelCase__      :			Optional[Any], UpperCamelCase__      :			int, UpperCamelCase__      :			Optional[Any]     )  ->			Any:
        if not self.initialized:
             _A    		=  RagRetriever(
                 UpperCamelCase__, question_encoder_tokenizer=UpperCamelCase__, generator_tokenizer=UpperCamelCase__, index=UpperCamelCase__, init_retrieval=UpperCamelCase__, )
             _A    		=  True
   def 	__UpperCAmelCase   (      self      :			Any     )  ->			List[Any]:
        self.retriever.index.init_index()

   def 	__UpperCAmelCase   (      self      :			int, UpperCamelCase__      :			Optional[Any], UpperCamelCase__      :			Dict     )  ->			Optional[Any]:
        _A		,       _A    		=  self.retriever._main_retrieve(UpperCamelCase__, UpperCamelCase__     )
        return doc_ids, retrieved_doc_embeds







class lowercase_       (		_UpperCamelCase			):







   """simple docstring"""

   def __init__(      self      :			Optional[Any], UpperCamelCase__      :			Optional[Any], UpperCamelCase__      :			Optional[int], UpperCamelCase__      :			Dict, UpperCamelCase__      :			Union[str, Any], UpperCamelCase__      :			int=None     )  ->			Optional[int]:
        if index is not None and index.is_initialized() and len(UpperCamelCase__     ) > 0:
             raise ValueError(
                 'When using Ray for distributed fine-tuning, '
                 'you\'ll need to provide the paths instead, '
                 'as the dataset and the index are loaded '
                 'separately. More info in examples/rag/use_own_knowledge_dataset.py '     )
        super().__init__(
            UpperCamelCase__, question_encoder_tokenizer=UpperCamelCase__, generator_tokenizer=UpperCamelCase__, index=UpperCamelCase__, init_retrieval=UpperCamelCase__, )
        _A    		=  retrieval_workers
        if len(self.retrieval_workers     ) > 0:
             ray.get(
                 [
                     worker.create_rag_retriever.remote(UpperCamelCase__, UpperCamelCase__, UpperCamelCase__, UpperCamelCase__     )
                     for worker in self.retrieval_workers
                 ]     )
   def 	__UpperCAmelCase   (      self      :			Dict     )  ->			int:
        logger.info('initializing retrieval'     )

        if len(self.retrieval_workers     ) > 0:
             ray.get([worker.init_retrieval.remote() for worker in self.retrieval_workers]     )
        else:
             # Non-distributed training. Load index into this same process.
             self.index.init_index()
   def 	__UpperCAmelCase   (      self      :			Optional[Any], UpperCamelCase__      :			Dict, UpperCamelCase__      :			Tuple     )  ->			Dict:
        if len(self.retrieval_workers     ) > 0:
             # Select a random retrieval actor.
             _A    		=  self.retrieval_workers[random.randint(0, len(self.retrieval_workers     ) - 1     )]
             _A		,       _A    		=  ray.get(random_worker.retrieve.remote(UpperCamelCase__, UpperCamelCase__     )     )
        else:
             _A		,       _A    		=  self._main_retrieve(UpperCamelCase__, UpperCamelCase__     )
        return retrieved_doc_embeds, doc_ids, self.index.get_doc_dicts(UpperCamelCase__     )
   @classmethod
   def 	__UpperCAmelCase   (      cls      :			List[str], UpperCamelCase__      :			List[Any], UpperCamelCase__      :			List[Any]=None, **UpperCamelCase__      :			Optional[int]     )  ->			Optional[Any]:
        return super(UpperCamelCase__, cls     ).get_tokenizers(UpperCamelCase__, UpperCamelCase__, **UpperCamelCase__     )

   @classmethod
   def 	__UpperCAmelCase   (      cls      :			Any, UpperCamelCase__      :			Tuple, UpperCamelCase__      :			int, UpperCamelCase__      :			List[Any]=None, **UpperCamelCase__      :			List[str]     )  ->			Tuple:
        _A    		=  kwargs.pop('config', UpperCamelCase__     ) or RagConfig.from_pretrained(UpperCamelCase__, **UpperCamelCase__     )
        _A    		=  RagTokenizer.from_pretrained(UpperCamelCase__, config=UpperCamelCase__     )
        _A    		=  rag_tokenizer.question_encoder
        _A    		=  rag_tokenizer.generator
        if indexed_dataset is not None:
             _A    		=  'custom'
             _A    		=  CustomHFIndex(config.retrieval_vector_size, UpperCamelCase__     )
        else:
             _A    		=  cls._build_index(UpperCamelCase__     )
        return cls(
            UpperCamelCase__, question_encoder_tokenizer=UpperCamelCase__, generator_tokenizer=UpperCamelCase__, retrieval_workers=UpperCamelCase__, index=UpperCamelCase__, )



