




"""simple docstring"""


def       UpperCAmelCase__ (snake_case__							:      int = 10_00					):
  """simple docstring"""

  return sum(e for e in range(3	,     snake_case__					) if e % 3 == 0 or e % 5 == 0					)


if __name__ == "__main__":
     print(F'''{solution() = }''')








"""simple docstring"""


import argparse
import json

import torch

from diffusers import DDPMScheduler, LDMPipeline, UNetaDModel, VQModel


def       UpperCAmelCase__ (snake_case__							:      Optional[Any]	,     snake_case__							:      Union[str, Any]=1					):
  """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__ (snake_case__							:      str	,     snake_case__							:      List[Any]=0					):
  """simple docstring"""

  _snake_case							:	Optional[Any]			=						[]
  for old_item in old_list:
    _snake_case							:	Union[str, Any]			=						old_item.replace("""in_layers.0"""	,     """norm1"""					)
    _snake_case							:	List[Any]			=						new_item.replace("""in_layers.2"""	,     """conv1"""					)

    _snake_case							:	Tuple			=						new_item.replace("""out_layers.0"""	,     """norm2"""					)
    _snake_case							:	Dict			=						new_item.replace("""out_layers.3"""	,     """conv2"""					)

    _snake_case							:	int			=						new_item.replace("""emb_layers.1"""	,     """time_emb_proj"""					)
    _snake_case							:	Optional[Any]			=						new_item.replace("""skip_connection"""	,     """conv_shortcut"""					)

    _snake_case							:	str			=						shave_segments(snake_case__	,     n_shave_prefix_segments=snake_case__					)

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

  return mapping


def       UpperCAmelCase__ (snake_case__							:      Dict	,     snake_case__							:      Dict=0					):
  """simple docstring"""

  _snake_case							:	Dict			=						[]
  for old_item in old_list:
    _snake_case							:	Dict			=						old_item

    _snake_case							:	int			=						new_item.replace("""norm.weight"""	,     """group_norm.weight"""					)
    _snake_case							:	str			=						new_item.replace("""norm.bias"""	,     """group_norm.bias"""					)

    _snake_case							:	Optional[Any]			=						new_item.replace("""proj_out.weight"""	,     """proj_attn.weight"""					)
    _snake_case							:	Optional[Any]			=						new_item.replace("""proj_out.bias"""	,     """proj_attn.bias"""					)

    _snake_case							:	Optional[Any]			=						shave_segments(snake_case__	,     n_shave_prefix_segments=snake_case__					)

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

  return mapping


def       UpperCAmelCase__ (snake_case__							:      str	,     snake_case__							:      Union[str, Any]	,     snake_case__							:      List[str]	,     snake_case__							:      str=None	,     snake_case__							:      str=None	,     snake_case__							:      List[str]=None					):
  """simple docstring"""

  assert isinstance(snake_case__	,     snake_case__					), "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							:	Union[str, Any]			=						old_checkpoint[path]
      _snake_case							:	Optional[int]			=						old_tensor.shape[0] // 3

      _snake_case							:	List[Any]			=						(-1, channels) if len(old_tensor.shape					) == 3 else (-1)

      _snake_case							:	Union[str, Any]			=						old_tensor.shape[0] // config["""num_head_channels"""] // 3

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

      _snake_case							:	Union[str, Any]			=						query.reshape(snake_case__					)
      _snake_case							:	Tuple			=						key.reshape(snake_case__					)
      _snake_case							:	Any			=						value.reshape(snake_case__					)

  for path in paths:
    _snake_case							:	List[Any]			=						path["""new"""]

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

    # Global renaming happens here
    _snake_case							:	Union[str, Any]			=						new_path.replace("""middle_block.0"""	,     """mid_block.resnets.0"""					)
    _snake_case							:	str			=						new_path.replace("""middle_block.1"""	,     """mid_block.attentions.0"""					)
    _snake_case							:	Any			=						new_path.replace("""middle_block.2"""	,     """mid_block.resnets.1"""					)

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

        # proj_attn.weight has to be converted from conv 1D to linear
    if "proj_attn.weight" in new_path:
      _snake_case							:	Dict			=						old_checkpoint[path["""old"""]][:, :, 0]
    else:
      _snake_case							:	Optional[Any]			=						old_checkpoint[path["""old"""]]


def       UpperCAmelCase__ (snake_case__							:      Any	,     snake_case__							:      List[str]					):
  """simple docstring"""

  _snake_case							:	int			=						{}

  _snake_case							:	Tuple			=						checkpoint["""time_embed.0.weight"""]
  _snake_case							:	List[str]			=						checkpoint["""time_embed.0.bias"""]
  _snake_case							:	List[str]			=						checkpoint["""time_embed.2.weight"""]
  _snake_case							:	Tuple			=						checkpoint["""time_embed.2.bias"""]

  _snake_case							:	Dict			=						checkpoint["""input_blocks.0.0.weight"""]
  _snake_case							:	List[Any]			=						checkpoint["""input_blocks.0.0.bias"""]

  _snake_case							:	List[Any]			=						checkpoint["""out.0.weight"""]
  _snake_case							:	Any			=						checkpoint["""out.0.bias"""]
  _snake_case							:	Any			=						checkpoint["""out.2.weight"""]
  _snake_case							:	List[str]			=						checkpoint["""out.2.bias"""]

  # Retrieves the keys for the input blocks only
  _snake_case							:	List[str]			=						len({""".""".join(layer.split("""."""					)[:2]					) for layer in checkpoint if """input_blocks""" in layer}					)
  _snake_case							:	Any			=						{
      layer_id: [key for key in checkpoint if F"input_blocks.{layer_id}" in key]
      for layer_id in range(snake_case__					)
  }

  # Retrieves the keys for the middle blocks only
  _snake_case							:	Optional[int]			=						len({""".""".join(layer.split("""."""					)[:2]					) for layer in checkpoint if """middle_block""" in layer}					)
  _snake_case							:	Optional[int]			=						{
      layer_id: [key for key in checkpoint if F"middle_block.{layer_id}" in key]
      for layer_id in range(snake_case__					)
  }

  # Retrieves the keys for the output blocks only
  _snake_case							:	str			=						len({""".""".join(layer.split("""."""					)[:2]					) for layer in checkpoint if """output_blocks""" in layer}					)
  _snake_case							:	List[Any]			=						{
      layer_id: [key for key in checkpoint if F"output_blocks.{layer_id}" in key]
      for layer_id in range(snake_case__					)
  }

  for i in range(1	,     snake_case__					):
    _snake_case							:	Union[str, Any]			=						(i - 1) // (config["""num_res_blocks"""] + 1)
    _snake_case							:	int			=						(i - 1) % (config["""num_res_blocks"""] + 1)

    _snake_case							:	List[str]			=						[key for key in input_blocks[i] if F"input_blocks.{i}.0" in key]
    _snake_case							:	str			=						[key for key in input_blocks[i] if F"input_blocks.{i}.1" in key]

    if F"input_blocks.{i}.0.op.weight" in checkpoint:
      _snake_case							:	Union[str, Any]			=						checkpoint[
          F"input_blocks.{i}.0.op.weight"
      ]
      _snake_case							:	Dict			=						checkpoint[
          F"input_blocks.{i}.0.op.bias"
      ]
      continue

    _snake_case							:	Optional[int]			=						renew_resnet_paths(snake_case__					)
    _snake_case							:	int			=						{"""old""": F"input_blocks.{i}.0", """new""": F"down_blocks.{block_id}.resnets.{layer_in_block_id}"}
    _snake_case							:	Tuple			=						{"""old""": """resnets.2.op""", """new""": """downsamplers.0.op"""}
    assign_to_checkpoint(
        snake_case__	,     snake_case__	,     snake_case__	,     additional_replacements=[meta_path, resnet_op]	,     config=snake_case__					)

    if len(snake_case__					):
      _snake_case							:	str			=						renew_attention_paths(snake_case__					)
      _snake_case							:	List[str]			=						{
          """old""": F"input_blocks.{i}.1",
          """new""": F"down_blocks.{block_id}.attentions.{layer_in_block_id}",
      }
      _snake_case							:	Optional[int]			=						{
          F"input_blocks.{i}.1.qkv.bias": {
              """key""": F"down_blocks.{block_id}.attentions.{layer_in_block_id}.key.bias",
              """query""": F"down_blocks.{block_id}.attentions.{layer_in_block_id}.query.bias",
              """value""": F"down_blocks.{block_id}.attentions.{layer_in_block_id}.value.bias",
          },
          F"input_blocks.{i}.1.qkv.weight": {
              """key""": F"down_blocks.{block_id}.attentions.{layer_in_block_id}.key.weight",
              """query""": F"down_blocks.{block_id}.attentions.{layer_in_block_id}.query.weight",
              """value""": F"down_blocks.{block_id}.attentions.{layer_in_block_id}.value.weight",
          },
      }
      assign_to_checkpoint(
          snake_case__	,     snake_case__	,     snake_case__	,     additional_replacements=[meta_path]	,     attention_paths_to_split=snake_case__	,     config=snake_case__	,     )

  _snake_case							:	int			=						middle_blocks[0]
  _snake_case							:	List[str]			=						middle_blocks[1]
  _snake_case							:	Any			=						middle_blocks[2]

  _snake_case							:	Dict			=						renew_resnet_paths(snake_case__					)
  assign_to_checkpoint(snake_case__	,     snake_case__	,     snake_case__	,     config=snake_case__					)

  _snake_case							:	Any			=						renew_resnet_paths(snake_case__					)
  assign_to_checkpoint(snake_case__	,     snake_case__	,     snake_case__	,     config=snake_case__					)

  _snake_case							:	Dict			=						renew_attention_paths(snake_case__					)
  _snake_case							:	Tuple			=						{
      """middle_block.1.qkv.bias""": {
          """key""": """mid_block.attentions.0.key.bias""",
          """query""": """mid_block.attentions.0.query.bias""",
          """value""": """mid_block.attentions.0.value.bias""",
      },
      """middle_block.1.qkv.weight""": {
          """key""": """mid_block.attentions.0.key.weight""",
          """query""": """mid_block.attentions.0.query.weight""",
          """value""": """mid_block.attentions.0.value.weight""",
      },
  }
  assign_to_checkpoint(
      snake_case__	,     snake_case__	,     snake_case__	,     attention_paths_to_split=snake_case__	,     config=snake_case__					)

  for i in range(snake_case__					):
    _snake_case							:	Optional[Any]			=						i // (config["""num_res_blocks"""] + 1)
    _snake_case							:	Dict			=						i % (config["""num_res_blocks"""] + 1)
    _snake_case							:	List[str]			=						[shave_segments(snake_case__	,     2					) for name in output_blocks[i]]
    _snake_case							:	Any			=						{}

    for layer in output_block_layers:
      _snake_case	,      _snake_case							:	Any			=						layer.split("""."""					)[0], shave_segments(snake_case__	,     1					)
      if layer_id in output_block_list:
        output_block_list[layer_id].append(snake_case__					)
      else:
        _snake_case							:	str			=						[layer_name]

    if len(snake_case__					) > 1:
      _snake_case							:	Dict			=						[key for key in output_blocks[i] if F"output_blocks.{i}.0" in key]
      _snake_case							:	List[str]			=						[key for key in output_blocks[i] if F"output_blocks.{i}.1" in key]

      _snake_case							:	List[Any]			=						renew_resnet_paths(snake_case__					)
      _snake_case							:	int			=						renew_resnet_paths(snake_case__					)

      _snake_case							:	Optional[Any]			=						{"""old""": F"output_blocks.{i}.0", """new""": F"up_blocks.{block_id}.resnets.{layer_in_block_id}"}
      assign_to_checkpoint(snake_case__	,     snake_case__	,     snake_case__	,     additional_replacements=[meta_path]	,     config=snake_case__					)

      if ["conv.weight", "conv.bias"] in output_block_list.values():
        _snake_case							:	str			=						list(output_block_list.values()					).index(["""conv.weight""", """conv.bias"""]					)
        _snake_case							:	Any			=						checkpoint[
            F"output_blocks.{i}.{index}.conv.weight"
        ]
        _snake_case							:	Optional[int]			=						checkpoint[
            F"output_blocks.{i}.{index}.conv.bias"
        ]

        # Clear attentions as they have been attributed above.
        if len(snake_case__					) == 2:
          _snake_case							:	Any			=						[]

      if len(snake_case__					):
        _snake_case							:	str			=						renew_attention_paths(snake_case__					)
        _snake_case							:	str			=						{
            """old""": F"output_blocks.{i}.1",
            """new""": F"up_blocks.{block_id}.attentions.{layer_in_block_id}",
        }
        _snake_case							:	int			=						{
            F"output_blocks.{i}.1.qkv.bias": {
                """key""": F"up_blocks.{block_id}.attentions.{layer_in_block_id}.key.bias",
                """query""": F"up_blocks.{block_id}.attentions.{layer_in_block_id}.query.bias",
                """value""": F"up_blocks.{block_id}.attentions.{layer_in_block_id}.value.bias",
            },
            F"output_blocks.{i}.1.qkv.weight": {
                """key""": F"up_blocks.{block_id}.attentions.{layer_in_block_id}.key.weight",
                """query""": F"up_blocks.{block_id}.attentions.{layer_in_block_id}.query.weight",
                """value""": F"up_blocks.{block_id}.attentions.{layer_in_block_id}.value.weight",
            },
        }
        assign_to_checkpoint(
            snake_case__	,     snake_case__	,     snake_case__	,     additional_replacements=[meta_path]	,     attention_paths_to_split=to_split if any("""qkv""" in key for key in attentions					) else None	,     config=snake_case__	,     )
    else:
      _snake_case							:	Optional[Any]			=						renew_resnet_paths(snake_case__	,     n_shave_prefix_segments=1					)
      for path in resnet_0_paths:
        _snake_case							:	Optional[Any]			=						""".""".join(["""output_blocks""", str(snake_case__					), path["""old"""]]					)
        _snake_case							:	Optional[int]			=						""".""".join(["""up_blocks""", str(snake_case__					), """resnets""", str(snake_case__					), path["""new"""]]					)

        _snake_case							:	Any			=						checkpoint[old_path]

  return new_checkpoint


if __name__ == "__main__":
     A_      			=       argparse.ArgumentParser()

     parser.add_argument(
         '''--checkpoint_path''', default=None, type=str, required=True, help='''Path to the checkpoint to convert.'''
     )

     parser.add_argument(
         '''--config_file''',
         default=None,
         type=str,
         required=True,
         help='''The config json file corresponding to the architecture.''',
     )

     parser.add_argument('''--dump_path''', default=None, type=str, required=True, help='''Path to the output model.''')

     A_      			=       parser.parse_args()

     A_      			=       torch.load(args.checkpoint_path)

     with open(args.config_file) as f:
          A_      			=       json.loads(f.read())

     A_      			=       convert_ldm_checkpoint(checkpoint, config)

     if "ldm" in config:
          del config["ldm"]

     A_      			=       UNetaDModel(**config)
     model.load_state_dict(converted_checkpoint)

     try:
          A_      			=       DDPMScheduler.from_config('''/'''.join(args.checkpoint_path.split('''/''')[:-1]))
          A_      			=       VQModel.from_pretrained('''/'''.join(args.checkpoint_path.split('''/''')[:-1]))

          A_      			=       LDMPipeline(unet=model, scheduler=scheduler, vae=vqvae)
          pipe.save_pretrained(args.dump_path)
     except:  # noqa: E722
          model.save_pretrained(args.dump_path)


