




'''simple docstring'''



import argparse

import torch

from diffusers import UNetaDConditionModel







def 	_UpperCamelCase       (_lowerCamelCase							:   int      ,		_lowerCamelCase							:   str      ,		_lowerCamelCase							:   Dict      ,		_lowerCamelCase							:   Optional[Any]=None      ,		_lowerCamelCase							:   Union[str, Any]=None      ,		_lowerCamelCase							:   Union[str, Any]=None		)->      Union[str, Any]:




     '''simple docstring'''





     assert isinstance(_lowerCamelCase      ,		_lowerCamelCase		), "Paths should be a list of dicts containing 'old' and 'new' keys."

     # Splits the attention layers into three variables.
     if attention_paths_to_split is not None:
          for path, path_map in attention_paths_to_split.items():
               __snake_case			    =     old_checkpoint[path]
               __snake_case			    =     old_tensor.shape[0] // 3

               __snake_case			    =     (-1, channels) if len(old_tensor.shape		) == 3 else (-1)

               __snake_case			    =     old_tensor.shape[0] // config['''num_head_channels'''] // 3

               __snake_case			    =     old_tensor.reshape((num_heads, 3 * channels // num_heads) + old_tensor.shape[1:]		)
               __snake_case    ,			__snake_case    ,			__snake_case			    =     old_tensor.split(channels // num_heads      ,		dim=1		)

               __snake_case			    =     query.reshape(_lowerCamelCase		)
               __snake_case			    =     key.reshape(_lowerCamelCase		)
               __snake_case			    =     value.reshape(_lowerCamelCase		)

     for path in paths:
          __snake_case			    =     path['''new''']

          # These have already been assigned
          if attention_paths_to_split is not None and new_path in attention_paths_to_split:
               continue

          if additional_replacements is not None:
               for replacement in additional_replacements:
                    __snake_case			    =     new_path.replace(replacement['''old''']      ,		replacement['''new''']		)

        # proj_attn.weight has to be converted from conv 1D to linear
          __snake_case			    =     old_checkpoint[path['''old''']]
          __snake_case			    =     ['''proj_attn.weight''']
          __snake_case			    =     ['''proj_out.weight''', '''proj_in.weight''']
          if any(k in new_path for k in names		):
               __snake_case			    =     weight[:, :, 0]
          elif any(k in new_path for k in names_a		) and len(weight.shape		) > 2 and ".attentions." not in new_path:
               __snake_case			    =     weight[:, :, 0]
          else:
               __snake_case			    =     weight







def 	_UpperCamelCase       (_lowerCamelCase							:   Tuple      ,		_lowerCamelCase							:   Union[str, Any]=0		)->      Union[str, Any]:




     '''simple docstring'''





     __snake_case			    =     []
     for old_item in old_list:
          __snake_case			    =     old_item

          #         new_item = new_item.replace('norm.weight', 'group_norm.weight')
          #         new_item = new_item.replace('norm.bias', 'group_norm.bias')

          #         new_item = new_item.replace('proj_out.weight', 'proj_attn.weight')
          #         new_item = new_item.replace('proj_out.bias', 'proj_attn.bias')

          #         new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments)

          mapping.append({'''old''': old_item, '''new''': new_item}		)

     return mapping







def 	_UpperCamelCase       (_lowerCamelCase							:   Any      ,		_lowerCamelCase							:   str=1		)->      Dict:




     '''simple docstring'''





     if n_shave_prefix_segments >= 0:
          return ".".join(path.split('''.'''		)[n_shave_prefix_segments:]		)
     else:
          return ".".join(path.split('''.'''		)[:n_shave_prefix_segments]		)







def 	_UpperCamelCase       (_lowerCamelCase							:   List[Any]      ,		_lowerCamelCase							:   Any=0		)->      Optional[Any]:




     '''simple docstring'''





     __snake_case			    =     []
     for old_item in old_list:
          mapping.append({'''old''': old_item, '''new''': old_item}		)

     return mapping







def 	_UpperCamelCase       (_lowerCamelCase							:   Any      ,		_lowerCamelCase							:   str=0		)->      Dict:




     '''simple docstring'''





     __snake_case			    =     []
     for old_item in old_list:
          __snake_case			    =     old_item.replace('''in_layers.0'''      ,		'''norm1'''		)
          __snake_case			    =     new_item.replace('''in_layers.2'''      ,		'''conv1'''		)

          __snake_case			    =     new_item.replace('''out_layers.0'''      ,		'''norm2'''		)
          __snake_case			    =     new_item.replace('''out_layers.3'''      ,		'''conv2'''		)

          __snake_case			    =     new_item.replace('''emb_layers.1'''      ,		'''time_emb_proj'''		)
          __snake_case			    =     new_item.replace('''skip_connection'''      ,		'''conv_shortcut'''		)

          __snake_case			    =     shave_segments(_lowerCamelCase      ,		n_shave_prefix_segments=_lowerCamelCase		)

          if "temopral_conv" not in old_item:
               mapping.append({'''old''': old_item, '''new''': new_item}		)

     return mapping







def 	_UpperCamelCase       (_lowerCamelCase							:   str      ,		_lowerCamelCase							:   Union[str, Any]      ,		_lowerCamelCase							:   int=None      ,		_lowerCamelCase							:   str=False		)->      Optional[int]:




     '''simple docstring'''





     __snake_case			    =     {}
     __snake_case			    =     list(checkpoint.keys()		)

     __snake_case			    =     '''model.diffusion_model.'''

     # at least a 100 parameters have to start with `model_ema` in order for the checkpoint to be EMA
     if sum(k.startswith('''model_ema'''		) for k in keys		) > 1_00 and extract_ema:
          print(f'''Checkpoint {path} has both EMA and non-EMA weights.'''		)
          print(
              '''In this conversion only the EMA weights are extracted. If you want to instead extract the non-EMA'''
              ''' weights (useful to continue fine-tuning), please make sure to remove the `--extract_ema` flag.'''		)
          for key in keys:
               if key.startswith('''model.diffusion_model'''		):
                    __snake_case			    =     '''model_ema.''' + ''''''.join(key.split('''.'''		)[1:]		)
                    __snake_case			    =     checkpoint.pop(_lowerCamelCase		)
     else:
          if sum(k.startswith('''model_ema'''		) for k in keys		) > 1_00:
               print(
                   '''In this conversion only the non-EMA weights are extracted. If you want to instead extract the EMA'''
                   ''' weights (usually better for inference), please make sure to add the `--extract_ema` flag.'''		)

          for key in keys:
               __snake_case			    =     checkpoint.pop(_lowerCamelCase		)

     __snake_case			    =     {}

     __snake_case			    =     unet_state_dict['''time_embed.0.weight''']
     __snake_case			    =     unet_state_dict['''time_embed.0.bias''']
     __snake_case			    =     unet_state_dict['''time_embed.2.weight''']
     __snake_case			    =     unet_state_dict['''time_embed.2.bias''']

     if config["class_embed_type"] is None:
          # No parameters to port
          ...
     elif config["class_embed_type"] == "timestep" or config["class_embed_type"] == "projection":
          __snake_case			    =     unet_state_dict['''label_emb.0.0.weight''']
          __snake_case			    =     unet_state_dict['''label_emb.0.0.bias''']
          __snake_case			    =     unet_state_dict['''label_emb.0.2.weight''']
          __snake_case			    =     unet_state_dict['''label_emb.0.2.bias''']
     else:
          raise NotImplementedError(f'''Not implemented `class_embed_type`: {config["class_embed_type"]}'''		)

     __snake_case			    =     unet_state_dict['''input_blocks.0.0.weight''']
     __snake_case			    =     unet_state_dict['''input_blocks.0.0.bias''']

     __snake_case			    =     [v for v in unet_state_dict if v.startswith('''input_blocks.0.1'''		)]
     __snake_case			    =     renew_attention_paths(_lowerCamelCase		)
     __snake_case			    =     {'''old''': '''input_blocks.0.1''', '''new''': '''transformer_in'''}
     assign_to_checkpoint(_lowerCamelCase      ,		_lowerCamelCase      ,		_lowerCamelCase      ,		additional_replacements=[meta_path]      ,		config=_lowerCamelCase		)

     __snake_case			    =     unet_state_dict['''out.0.weight''']
     __snake_case			    =     unet_state_dict['''out.0.bias''']
     __snake_case			    =     unet_state_dict['''out.2.weight''']
     __snake_case			    =     unet_state_dict['''out.2.bias''']

     # Retrieves the keys for the input blocks only
     __snake_case			    =     len({'''.'''.join(layer.split('''.'''		)[:2]		) for layer in unet_state_dict if '''input_blocks''' in layer}		)
     __snake_case			    =     {
         layer_id: [key for key in unet_state_dict if f'''input_blocks.{layer_id}''' in key]
         for layer_id in range(_lowerCamelCase		)
     }

     # Retrieves the keys for the middle blocks only
     __snake_case			    =     len({'''.'''.join(layer.split('''.'''		)[:2]		) for layer in unet_state_dict if '''middle_block''' in layer}		)
     __snake_case			    =     {
         layer_id: [key for key in unet_state_dict if f'''middle_block.{layer_id}''' in key]
         for layer_id in range(_lowerCamelCase		)
     }

     # Retrieves the keys for the output blocks only
     __snake_case			    =     len({'''.'''.join(layer.split('''.'''		)[:2]		) for layer in unet_state_dict if '''output_blocks''' in layer}		)
     __snake_case			    =     {
         layer_id: [key for key in unet_state_dict if f'''output_blocks.{layer_id}''' in key]
         for layer_id in range(_lowerCamelCase		)
     }

     for i in range(1      ,		_lowerCamelCase		):
          __snake_case			    =     (i - 1) // (config['''layers_per_block'''] + 1)
          __snake_case			    =     (i - 1) % (config['''layers_per_block'''] + 1)

          __snake_case			    =     [
              key for key in input_blocks[i] if f'''input_blocks.{i}.0''' in key and f'''input_blocks.{i}.0.op''' not in key
          ]
          __snake_case			    =     [key for key in input_blocks[i] if f'''input_blocks.{i}.1''' in key]
          __snake_case			    =     [key for key in input_blocks[i] if f'''input_blocks.{i}.2''' in key]

          if f'''input_blocks.{i}.op.weight''' in unet_state_dict:
               __snake_case			    =     unet_state_dict.pop(
                   f'''input_blocks.{i}.op.weight'''		)
               __snake_case			    =     unet_state_dict.pop(
                   f'''input_blocks.{i}.op.bias'''		)

          __snake_case			    =     renew_resnet_paths(_lowerCamelCase		)
          __snake_case			    =     {'''old''': f'''input_blocks.{i}.0''', '''new''': f'''down_blocks.{block_id}.resnets.{layer_in_block_id}'''}
          assign_to_checkpoint(
              _lowerCamelCase      ,		_lowerCamelCase      ,		_lowerCamelCase      ,		additional_replacements=[meta_path]      ,		config=_lowerCamelCase		)

          __snake_case			    =     [key for key in resnets if '''temopral_conv''' in key]
          __snake_case			    =     renew_temp_conv_paths(_lowerCamelCase		)
          __snake_case			    =     {
              '''old''': f'''input_blocks.{i}.0.temopral_conv''',
              '''new''': f'''down_blocks.{block_id}.temp_convs.{layer_in_block_id}''',
          }
          assign_to_checkpoint(
              _lowerCamelCase      ,		_lowerCamelCase      ,		_lowerCamelCase      ,		additional_replacements=[meta_path]      ,		config=_lowerCamelCase		)

          if len(_lowerCamelCase		):
               __snake_case			    =     renew_attention_paths(_lowerCamelCase		)
               __snake_case			    =     {'''old''': f'''input_blocks.{i}.1''', '''new''': f'''down_blocks.{block_id}.attentions.{layer_in_block_id}'''}
               assign_to_checkpoint(
                   _lowerCamelCase      ,		_lowerCamelCase      ,		_lowerCamelCase      ,		additional_replacements=[meta_path]      ,		config=_lowerCamelCase		)

          if len(_lowerCamelCase		):
               __snake_case			    =     renew_attention_paths(_lowerCamelCase		)
               __snake_case			    =     {
                   '''old''': f'''input_blocks.{i}.2''',
                   '''new''': f'''down_blocks.{block_id}.temp_attentions.{layer_in_block_id}''',
               }
               assign_to_checkpoint(
                   _lowerCamelCase      ,		_lowerCamelCase      ,		_lowerCamelCase      ,		additional_replacements=[meta_path]      ,		config=_lowerCamelCase		)

     __snake_case			    =     middle_blocks[0]
     __snake_case			    =     [key for key in resnet_a if '''temopral_conv''' in key]
     __snake_case			    =     middle_blocks[1]
     __snake_case			    =     middle_blocks[2]
     __snake_case			    =     middle_blocks[3]
     __snake_case			    =     [key for key in resnet_a if '''temopral_conv''' in key]

     __snake_case			    =     renew_resnet_paths(_lowerCamelCase		)
     __snake_case			    =     {'''old''': '''middle_block.0''', '''new''': '''mid_block.resnets.0'''}
     assign_to_checkpoint(
         _lowerCamelCase      ,		_lowerCamelCase      ,		_lowerCamelCase      ,		config=_lowerCamelCase      ,		additional_replacements=[meta_path]		)

     __snake_case			    =     renew_temp_conv_paths(_lowerCamelCase		)
     __snake_case			    =     {'''old''': '''middle_block.0.temopral_conv''', '''new''': '''mid_block.temp_convs.0'''}
     assign_to_checkpoint(
         _lowerCamelCase      ,		_lowerCamelCase      ,		_lowerCamelCase      ,		config=_lowerCamelCase      ,		additional_replacements=[meta_path]		)

     __snake_case			    =     renew_resnet_paths(_lowerCamelCase		)
     __snake_case			    =     {'''old''': '''middle_block.3''', '''new''': '''mid_block.resnets.1'''}
     assign_to_checkpoint(
         _lowerCamelCase      ,		_lowerCamelCase      ,		_lowerCamelCase      ,		config=_lowerCamelCase      ,		additional_replacements=[meta_path]		)

     __snake_case			    =     renew_temp_conv_paths(_lowerCamelCase		)
     __snake_case			    =     {'''old''': '''middle_block.3.temopral_conv''', '''new''': '''mid_block.temp_convs.1'''}
     assign_to_checkpoint(
         _lowerCamelCase      ,		_lowerCamelCase      ,		_lowerCamelCase      ,		config=_lowerCamelCase      ,		additional_replacements=[meta_path]		)

     __snake_case			    =     renew_attention_paths(_lowerCamelCase		)
     __snake_case			    =     {'''old''': '''middle_block.1''', '''new''': '''mid_block.attentions.0'''}
     assign_to_checkpoint(
         _lowerCamelCase      ,		_lowerCamelCase      ,		_lowerCamelCase      ,		additional_replacements=[meta_path]      ,		config=_lowerCamelCase		)

     __snake_case			    =     renew_attention_paths(_lowerCamelCase		)
     __snake_case			    =     {'''old''': '''middle_block.2''', '''new''': '''mid_block.temp_attentions.0'''}
     assign_to_checkpoint(
         _lowerCamelCase      ,		_lowerCamelCase      ,		_lowerCamelCase      ,		additional_replacements=[meta_path]      ,		config=_lowerCamelCase		)

     for i in range(_lowerCamelCase		):
          __snake_case			    =     i // (config['''layers_per_block'''] + 1)
          __snake_case			    =     i % (config['''layers_per_block'''] + 1)
          __snake_case			    =     [shave_segments(_lowerCamelCase      ,		2		) for name in output_blocks[i]]
          __snake_case			    =     {}

          for layer in output_block_layers:
               __snake_case    ,			__snake_case			    =     layer.split('''.'''		)[0], shave_segments(_lowerCamelCase      ,		1		)
               if layer_id in output_block_list:
                    output_block_list[layer_id].append(_lowerCamelCase		)
               else:
                    __snake_case			    =     [layer_name]

          if len(_lowerCamelCase		) > 1:
               __snake_case			    =     [key for key in output_blocks[i] if f'''output_blocks.{i}.0''' in key]
               __snake_case			    =     [key for key in output_blocks[i] if f'''output_blocks.{i}.1''' in key]
               __snake_case			    =     [key for key in output_blocks[i] if f'''output_blocks.{i}.2''' in key]

               __snake_case			    =     renew_resnet_paths(_lowerCamelCase		)
               __snake_case			    =     renew_resnet_paths(_lowerCamelCase		)

               __snake_case			    =     {'''old''': f'''output_blocks.{i}.0''', '''new''': f'''up_blocks.{block_id}.resnets.{layer_in_block_id}'''}
               assign_to_checkpoint(
                   _lowerCamelCase      ,		_lowerCamelCase      ,		_lowerCamelCase      ,		additional_replacements=[meta_path]      ,		config=_lowerCamelCase		)

               __snake_case			    =     [key for key in resnets if '''temopral_conv''' in key]
               __snake_case			    =     renew_temp_conv_paths(_lowerCamelCase		)
               __snake_case			    =     {
                   '''old''': f'''output_blocks.{i}.0.temopral_conv''',
                   '''new''': f'''up_blocks.{block_id}.temp_convs.{layer_in_block_id}''',
               }
               assign_to_checkpoint(
                   _lowerCamelCase      ,		_lowerCamelCase      ,		_lowerCamelCase      ,		additional_replacements=[meta_path]      ,		config=_lowerCamelCase		)

               __snake_case			    =     {k: sorted(_lowerCamelCase		) for k, v in output_block_list.items()}
               if ["conv.bias", "conv.weight"] in output_block_list.values():
                    __snake_case			    =     list(output_block_list.values()		).index(['''conv.bias''', '''conv.weight''']		)
                    __snake_case			    =     unet_state_dict[
                        f'''output_blocks.{i}.{index}.conv.weight'''
                    ]
                    __snake_case			    =     unet_state_dict[
                        f'''output_blocks.{i}.{index}.conv.bias'''
                    ]

                    # Clear attentions as they have been attributed above.
                    if len(_lowerCamelCase		) == 2:
                         __snake_case			    =     []

               if len(_lowerCamelCase		):
                    __snake_case			    =     renew_attention_paths(_lowerCamelCase		)
                    __snake_case			    =     {
                        '''old''': f'''output_blocks.{i}.1''',
                        '''new''': f'''up_blocks.{block_id}.attentions.{layer_in_block_id}''',
                    }
                    assign_to_checkpoint(
                        _lowerCamelCase      ,		_lowerCamelCase      ,		_lowerCamelCase      ,		additional_replacements=[meta_path]      ,		config=_lowerCamelCase		)

               if len(_lowerCamelCase		):
                    __snake_case			    =     renew_attention_paths(_lowerCamelCase		)
                    __snake_case			    =     {
                        '''old''': f'''output_blocks.{i}.2''',
                        '''new''': f'''up_blocks.{block_id}.temp_attentions.{layer_in_block_id}''',
                    }
                    assign_to_checkpoint(
                        _lowerCamelCase      ,		_lowerCamelCase      ,		_lowerCamelCase      ,		additional_replacements=[meta_path]      ,		config=_lowerCamelCase		)
          else:
               __snake_case			    =     renew_resnet_paths(_lowerCamelCase      ,		n_shave_prefix_segments=1		)
               for path in resnet_0_paths:
                    __snake_case			    =     '''.'''.join(['''output_blocks''', str(_lowerCamelCase		), path['''old''']]		)
                    __snake_case			    =     '''.'''.join(['''up_blocks''', str(_lowerCamelCase		), '''resnets''', str(_lowerCamelCase		), path['''new''']]		)
                    __snake_case			    =     unet_state_dict[old_path]

               __snake_case			    =     [l for l in output_block_layers if '''temopral_conv''' in l]
               for path in temopral_conv_paths:
                    __snake_case			    =     path.split('''temopral_conv.'''		)[-1]
                    __snake_case			    =     '''.'''.join(['''output_blocks''', str(_lowerCamelCase		), str(_lowerCamelCase		), '''temopral_conv''', pruned_path]		)
                    __snake_case			    =     '''.'''.join(['''up_blocks''', str(_lowerCamelCase		), '''temp_convs''', str(_lowerCamelCase		), pruned_path]		)
                    __snake_case			    =     unet_state_dict[old_path]

     return new_checkpoint


if __name__ == "__main__":
     UpperCAmelCase_				:     int     			=						argparse.ArgumentParser()

     parser.add_argument(
         '''--checkpoint_path''', default=None, type=str, required=True, help='''Path to the checkpoint to convert.'''
     )
     parser.add_argument('''--dump_path''', default=None, type=str, required=True, help='''Path to the output model.''')
     UpperCAmelCase_				:     int     			=						parser.parse_args()

     UpperCAmelCase_				:     str     			=						torch.load(args.checkpoint_path, map_location='''cpu''')
     UpperCAmelCase_				:     Any     			=						UNetaDConditionModel()

     UpperCAmelCase_				:     str     			=						convert_ldm_unet_checkpoint(unet_checkpoint, unet.config)

     UpperCAmelCase_				:     Union[str, Any]     			=						set(unet.state_dict().keys()) - set(converted_ckpt.keys())
     UpperCAmelCase_				:     Optional[Any]     			=						set(converted_ckpt.keys()) - set(unet.state_dict().keys())

     assert len(diff_a) == len(diff_a) == 0, "Converted weights don't match"

     # load state_dict
     unet.load_state_dict(converted_ckpt)

     unet.save_pretrained(args.dump_path)

     # -- finish converting the unet --












'''simple docstring'''



import os

from packaging import version

from .. import __version__
from .accelerate_utils import apply_forward_hook
from .constants import (
    CONFIG_NAME,
    DEPRECATED_REVISION_ARGS,
    DIFFUSERS_CACHE,
    DIFFUSERS_DYNAMIC_MODULE_NAME,
    FLAX_WEIGHTS_NAME,
    HF_MODULES_CACHE,
    HUGGINGFACE_CO_RESOLVE_ENDPOINT,
    ONNX_EXTERNAL_WEIGHTS_NAME,
    ONNX_WEIGHTS_NAME,
    SAFETENSORS_WEIGHTS_NAME,
    TEXT_ENCODER_ATTN_MODULE,
    WEIGHTS_NAME,
)
from .deprecation_utils import deprecate
from .doc_utils import replace_example_docstring
from .dynamic_modules_utils import get_class_from_dynamic_module
from .hub_utils import (
    HF_HUB_OFFLINE,
    _add_variant,
    _get_model_file,
    extract_commit_hash,
    http_user_agent,
)
from .import_utils import (
    BACKENDS_MAPPING,
    ENV_VARS_TRUE_AND_AUTO_VALUES,
    ENV_VARS_TRUE_VALUES,
    USE_JAX,
    USE_TF,
    USE_TORCH,
    DummyObject,
    OptionalDependencyNotAvailable,
    is_accelerate_available,
    is_accelerate_version,
    is_bsa_available,
    is_flax_available,
    is_ftfy_available,
    is_inflect_available,
    is_invisible_watermark_available,
    is_k_diffusion_available,
    is_k_diffusion_version,
    is_librosa_available,
    is_note_seq_available,
    is_omegaconf_available,
    is_onnx_available,
    is_safetensors_available,
    is_scipy_available,
    is_tensorboard_available,
    is_tf_available,
    is_torch_available,
    is_torch_version,
    is_torchsde_available,
    is_transformers_available,
    is_transformers_version,
    is_unidecode_available,
    is_wandb_available,
    is_xformers_available,
    requires_backends,
)
from .logging import get_logger
from .outputs import BaseOutput
from .pil_utils import PIL_INTERPOLATION, numpy_to_pil, pt_to_pil
from .torch_utils import is_compiled_module, randn_tensor


if is_torch_available():
     from .testing_utils import (
         floats_tensor,
         load_hf_numpy,
         load_image,
         load_numpy,
         load_pt,
         nightly,
         parse_flag_from_env,
         print_tensor_test,
         require_torch_a,
         require_torch_gpu,
         skip_mps,
         slow,
         torch_all_close,
         torch_device,
     )
     from .torch_utils import maybe_allow_in_graph

from .testing_utils import export_to_gif, export_to_video


UpperCAmelCase_				:     int     			=						get_logger(__name__)







def 	_UpperCamelCase       (_lowerCamelCase							:   Union[str, Any]		)->      Tuple:




     '''simple docstring'''





     if version.parse(_lowerCamelCase		) < version.parse(_lowerCamelCase		):
          if "dev" in min_version:
               __snake_case			    =     (
                   '''This example requires a source install from HuggingFace diffusers (see '''
                   '''`https://huggingface.co/docs/diffusers/installation#install-from-source`),'''
               )
          else:
               __snake_case			    =     f'''This example requires a minimum version of {min_version},'''
          error_message += f''' but the version found is {__version__}.\n'''
          raise ImportError(_lowerCamelCase		)






