


import argparse
import json

import torch

from diffusers import DDPMScheduler, LDMPipeline, UNetaDModel, VQModel
def 				UpperCamelCase__					(  SCREAMING_SNAKE_CASE__      ,	SCREAMING_SNAKE_CASE__=1						):
 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__					(  SCREAMING_SNAKE_CASE__      ,	SCREAMING_SNAKE_CASE__=0						):
 __lowerCamelCase		:							List[Any]		 =						[]
 for old_item in old_list:
  __lowerCamelCase		:							Union[str, Any]		 =						old_item.replace('in_layers.0'      ,	'norm1'						)
  __lowerCamelCase		:							Optional[int]		 =						new_item.replace('in_layers.2'      ,	'conv1'						)

  __lowerCamelCase		:							str		 =						new_item.replace('out_layers.0'      ,	'norm2'						)
  __lowerCamelCase		:							Union[str, Any]		 =						new_item.replace('out_layers.3'      ,	'conv2'						)

  __lowerCamelCase		:							Optional[int]		 =						new_item.replace('emb_layers.1'      ,	'time_emb_proj'						)
  __lowerCamelCase		:							List[str]		 =						new_item.replace('skip_connection'      ,	'conv_shortcut'						)

  __lowerCamelCase		:							Dict		 =						shave_segments(_UpperCAmelCase      ,	n_shave_prefix_segments=_UpperCAmelCase						)

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

 return mapping
def 				UpperCamelCase__					(  SCREAMING_SNAKE_CASE__      ,	SCREAMING_SNAKE_CASE__=0						):
 __lowerCamelCase		:							Union[str, Any]		 =						[]
 for old_item in old_list:
  __lowerCamelCase		:							Any		 =						old_item

  __lowerCamelCase		:							Tuple		 =						new_item.replace('norm.weight'      ,	'group_norm.weight'						)
  __lowerCamelCase		:							List[Any]		 =						new_item.replace('norm.bias'      ,	'group_norm.bias'						)

  __lowerCamelCase		:							Tuple		 =						new_item.replace('proj_out.weight'      ,	'proj_attn.weight'						)
  __lowerCamelCase		:							int		 =						new_item.replace('proj_out.bias'      ,	'proj_attn.bias'						)

  __lowerCamelCase		:							List[str]		 =						shave_segments(_UpperCAmelCase      ,	n_shave_prefix_segments=_UpperCAmelCase						)

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

 return mapping
def 				UpperCamelCase__					(  SCREAMING_SNAKE_CASE__      ,	SCREAMING_SNAKE_CASE__      ,	SCREAMING_SNAKE_CASE__      ,	SCREAMING_SNAKE_CASE__=None      ,	SCREAMING_SNAKE_CASE__=None      ,	SCREAMING_SNAKE_CASE__=None						):
 assert isinstance(_UpperCAmelCase      ,	_UpperCAmelCase						), "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():
   __lowerCamelCase		:							List[str]		 =						old_checkpoint[path]
   __lowerCamelCase		:							Dict		 =						old_tensor.shape[0] // 3

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

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

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

   __lowerCamelCase		:							List[str]		 =						query.reshape(_UpperCAmelCase						)
   __lowerCamelCase		:							List[Any]		 =						key.reshape(_UpperCAmelCase						)
   __lowerCamelCase		:							Union[str, Any]		 =						value.reshape(_UpperCAmelCase						)

 for path in paths:
  __lowerCamelCase		:							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
  __lowerCamelCase		:							Any		 =						new_path.replace('middle_block.0'      ,	'mid_block.resnets.0'						)
  __lowerCamelCase		:							int		 =						new_path.replace('middle_block.1'      ,	'mid_block.attentions.0'						)
  __lowerCamelCase		:							Optional[Any]		 =						new_path.replace('middle_block.2'      ,	'mid_block.resnets.1'						)

  if additional_replacements is not None:
   for replacement in additional_replacements:
    __lowerCamelCase		:							Optional[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:
   __lowerCamelCase		:							Tuple		 =						old_checkpoint[path["old"]][:, :, 0]
  else:
   __lowerCamelCase		:							int		 =						old_checkpoint[path["old"]]
def 				UpperCamelCase__					(  SCREAMING_SNAKE_CASE__      ,	SCREAMING_SNAKE_CASE__						):
 __lowerCamelCase		:							Any		 =						{}

 __lowerCamelCase		:							List[Any]		 =						checkpoint["time_embed.0.weight"]
 __lowerCamelCase		:							Optional[Any]		 =						checkpoint["time_embed.0.bias"]
 __lowerCamelCase		:							List[Any]		 =						checkpoint["time_embed.2.weight"]
 __lowerCamelCase		:							List[str]		 =						checkpoint["time_embed.2.bias"]

 __lowerCamelCase		:							Optional[Any]		 =						checkpoint["input_blocks.0.0.weight"]
 __lowerCamelCase		:							Any		 =						checkpoint["input_blocks.0.0.bias"]

 __lowerCamelCase		:							Optional[Any]		 =						checkpoint["out.0.weight"]
 __lowerCamelCase		:							List[Any]		 =						checkpoint["out.0.bias"]
 __lowerCamelCase		:							int		 =						checkpoint["out.2.weight"]
 __lowerCamelCase		:							List[Any]		 =						checkpoint["out.2.bias"]

 # Retrieves the keys for the input blocks only
 __lowerCamelCase		:							int		 =						len({'.'.join(layer.split('.'						)[:2]						) for layer in checkpoint if 'input_blocks' in layer}						)
 __lowerCamelCase		:							Union[str, Any]		 =						{
     layer_id: [key for key in checkpoint if f'input_blocks.{layer_id}' in key]
     for layer_id in range(_UpperCAmelCase						)
 }

 # Retrieves the keys for the middle blocks only
 __lowerCamelCase		:							Any		 =						len({'.'.join(layer.split('.'						)[:2]						) for layer in checkpoint if 'middle_block' in layer}						)
 __lowerCamelCase		:							Tuple		 =						{
     layer_id: [key for key in checkpoint if f'middle_block.{layer_id}' in key]
     for layer_id in range(_UpperCAmelCase						)
 }

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

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

  __lowerCamelCase		:							Any		 =						[key for key in input_blocks[i] if f'input_blocks.{i}.0' in key]
  __lowerCamelCase		:							Dict		 =						[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:
   __lowerCamelCase		:							Union[str, Any]		 =						checkpoint[
       f'input_blocks.{i}.0.op.weight'
   ]
   __lowerCamelCase		:							Tuple		 =						checkpoint[
       f'input_blocks.{i}.0.op.bias'
   ]
   continue

  __lowerCamelCase		:							Dict		 =						renew_resnet_paths(_UpperCAmelCase						)
  __lowerCamelCase		:							List[Any]		 =						{"old": f'input_blocks.{i}.0', "new": f'down_blocks.{block_id}.resnets.{layer_in_block_id}'}
  __lowerCamelCase		:							Union[str, Any]		 =						{"old": "resnets.2.op", "new": "downsamplers.0.op"}
  assign_to_checkpoint(
      _UpperCAmelCase      ,	_UpperCAmelCase      ,	_UpperCAmelCase      ,	additional_replacements=[meta_path, resnet_op]      ,	config=_UpperCAmelCase						)

  if len(_UpperCAmelCase						):
   __lowerCamelCase		:							Optional[int]		 =						renew_attention_paths(_UpperCAmelCase						)
   __lowerCamelCase		:							Optional[int]		 =						{
       "old": f'input_blocks.{i}.1',
       "new": f'down_blocks.{block_id}.attentions.{layer_in_block_id}',
   }
   __lowerCamelCase		:							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(
       _UpperCAmelCase      ,	_UpperCAmelCase      ,	_UpperCAmelCase      ,	additional_replacements=[meta_path]      ,	attention_paths_to_split=_UpperCAmelCase      ,	config=_UpperCAmelCase      ,	)

 __lowerCamelCase		:							Any		 =						middle_blocks[0]
 __lowerCamelCase		:							Optional[Any]		 =						middle_blocks[1]
 __lowerCamelCase		:							Optional[int]		 =						middle_blocks[2]

 __lowerCamelCase		:							List[Any]		 =						renew_resnet_paths(_UpperCAmelCase						)
 assign_to_checkpoint(_UpperCAmelCase      ,	_UpperCAmelCase      ,	_UpperCAmelCase      ,	config=_UpperCAmelCase						)

 __lowerCamelCase		:							Union[str, Any]		 =						renew_resnet_paths(_UpperCAmelCase						)
 assign_to_checkpoint(_UpperCAmelCase      ,	_UpperCAmelCase      ,	_UpperCAmelCase      ,	config=_UpperCAmelCase						)

 __lowerCamelCase		:							Optional[int]		 =						renew_attention_paths(_UpperCAmelCase						)
 __lowerCamelCase		:							Optional[Any]		 =						{
     "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(
     _UpperCAmelCase      ,	_UpperCAmelCase      ,	_UpperCAmelCase      ,	attention_paths_to_split=_UpperCAmelCase      ,	config=_UpperCAmelCase						)

 for i in range(_UpperCAmelCase						):
  __lowerCamelCase		:							Dict		 =						i // (config["num_res_blocks"] + 1)
  __lowerCamelCase		:							Any		 =						i % (config["num_res_blocks"] + 1)
  __lowerCamelCase		:							Tuple		 =						[shave_segments(_UpperCAmelCase      ,	2						) for name in output_blocks[i]]
  __lowerCamelCase		:							Optional[Any]		 =						{}

  for layer in output_block_layers:
   __lowerCamelCase		:							Optional[Any]		 =						layer.split('.'						)[0], shave_segments(_UpperCAmelCase      ,	1						)
   if layer_id in output_block_list:
    output_block_list[layer_id].append(_UpperCAmelCase						)
   else:
    __lowerCamelCase		:							List[str]		 =						[layer_name]

  if len(_UpperCAmelCase						) > 1:
   __lowerCamelCase		:							Union[str, Any]		 =						[key for key in output_blocks[i] if f'output_blocks.{i}.0' in key]
   __lowerCamelCase		:							Optional[int]		 =						[key for key in output_blocks[i] if f'output_blocks.{i}.1' in key]

   __lowerCamelCase		:							Tuple		 =						renew_resnet_paths(_UpperCAmelCase						)
   __lowerCamelCase		:							Tuple		 =						renew_resnet_paths(_UpperCAmelCase						)

   __lowerCamelCase		:							Optional[int]		 =						{"old": f'output_blocks.{i}.0', "new": f'up_blocks.{block_id}.resnets.{layer_in_block_id}'}
   assign_to_checkpoint(_UpperCAmelCase      ,	_UpperCAmelCase      ,	_UpperCAmelCase      ,	additional_replacements=[meta_path]      ,	config=_UpperCAmelCase						)

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

    # Clear attentions as they have been attributed above.
    if len(_UpperCAmelCase						) == 2:
     __lowerCamelCase		:							Union[str, Any]		 =						[]

   if len(_UpperCAmelCase						):
    __lowerCamelCase		:							List[str]		 =						renew_attention_paths(_UpperCAmelCase						)
    __lowerCamelCase		:							Any		 =						{
        "old": f'output_blocks.{i}.1',
        "new": f'up_blocks.{block_id}.attentions.{layer_in_block_id}',
    }
    __lowerCamelCase		:							Any		 =						{
        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(
        _UpperCAmelCase      ,	_UpperCAmelCase      ,	_UpperCAmelCase      ,	additional_replacements=[meta_path]      ,	attention_paths_to_split=to_split if any('qkv' in key for key in attentions						) else None      ,	config=_UpperCAmelCase      ,	)
  else:
   __lowerCamelCase		:							Tuple		 =						renew_resnet_paths(_UpperCAmelCase      ,	n_shave_prefix_segments=1						)
   for path in resnet_0_paths:
    __lowerCamelCase		:							Optional[int]		 =						".".join(['output_blocks', str(_UpperCAmelCase						), path['old']]						)
    __lowerCamelCase		:							Optional[int]		 =						".".join(['up_blocks', str(_UpperCAmelCase						), 'resnets', str(_UpperCAmelCase						), path['new']]						)

    __lowerCamelCase		:							List[Any]		 =						checkpoint[old_path]

 return new_checkpoint


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

       lowercase_	   =     parser.parse_args()

       lowercase_	   =     torch.load(args.checkpoint_path)

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

       lowercase_	   =     convert_ldm_checkpoint(checkpoint, config)

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

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

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

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





from math import pow
def 				UpperCamelCase__					(  SCREAMING_SNAKE_CASE__      ,	SCREAMING_SNAKE_CASE__      ,	SCREAMING_SNAKE_CASE__      ,	SCREAMING_SNAKE_CASE__      ,	SCREAMING_SNAKE_CASE__      ,	):
 if current_sum == needed_sum:
  # If the sum of the powers is equal to needed_sum, then we have a solution.
  solutions_count += 1
  return current_sum, solutions_count

 __lowerCamelCase		:							Optional[Any]		 =						int(pow(SCREAMING_SNAKE_CASE__      ,	SCREAMING_SNAKE_CASE__						)						)
 if current_sum + i_to_n <= needed_sum:
  # If the sum of the powers is less than needed_sum, then continue adding powers.
  current_sum += i_to_n
  __lowerCamelCase ,					__lowerCamelCase		:							Optional[Any]		 =						backtrack(
      SCREAMING_SNAKE_CASE__      ,	SCREAMING_SNAKE_CASE__      ,	current_number + 1      ,	SCREAMING_SNAKE_CASE__      ,	SCREAMING_SNAKE_CASE__						)
  current_sum -= i_to_n
 if i_to_n < needed_sum:
  # If the power of i is less than needed_sum, then try with the next power.
  __lowerCamelCase ,					__lowerCamelCase		:							Dict		 =						backtrack(
      SCREAMING_SNAKE_CASE__      ,	SCREAMING_SNAKE_CASE__      ,	current_number + 1      ,	SCREAMING_SNAKE_CASE__      ,	SCREAMING_SNAKE_CASE__						)
 return current_sum, solutions_count
def 				UpperCamelCase__					(  SCREAMING_SNAKE_CASE__      ,	SCREAMING_SNAKE_CASE__						):
 if not (1 <= needed_sum <= 1_000 and 2 <= power <= 10):
  raise ValueError(
      'Invalid input\n'
      'needed_sum must be between 1 and 1000, power between 2 and 10.'						)

 return backtrack(SCREAMING_SNAKE_CASE__      ,	SCREAMING_SNAKE_CASE__      ,	1      ,	0      ,	0						)[1]  # Return the solutions_count


if __name__ == "__main__":
       import doctest

       doctest.testmod()

