code: 
  
  
  
  
  
  import argparse
  
  import torch
  
  from diffusers import UNetaDConditionModel
  
  
  
  
  def 						A_			(  lowercase_      ,		lowercase_      ,		lowercase_      ,		lowercase_=None      ,		lowercase_=None      ,		lowercase_=None  )			->						Optional[Any]:
         assert isinstance(lowercase_      ,		lowercase_  ), "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      :							Optional[int]     				=						old_checkpoint[path]
                       _snake_case      :							Dict     				=						old_tensor.shape[0] // 3
  
                       _snake_case      :							List[Any]     				=						(-1, channels) if len(old_tensor.shape  ) == 3 else (-1)
  
                       _snake_case      :							Optional[int]     				=						old_tensor.shape[0] // config['''num_head_channels'''] // 3
  
                       _snake_case      :							Tuple     				=						old_tensor.reshape((num_heads, 3 * channels // num_heads) + old_tensor.shape[1:]  )
                       _snake_case				,    _snake_case				,    _snake_case      :							Union[str, Any]     				=						old_tensor.split(channels // num_heads      ,		dim=1  )
  
                       _snake_case      :							Optional[Any]     				=						query.reshape(lowercase_  )
                       _snake_case      :							Optional[Any]     				=						key.reshape(lowercase_  )
                       _snake_case      :							Tuple     				=						value.reshape(lowercase_  )
  
         for path in paths:
                _snake_case      :							Optional[int]     				=						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      :							Tuple     				=						new_path.replace(replacement['''old''']      ,		replacement['''new''']  )
  
          # proj_attn.weight has to be converted from conv 1D to linear
                _snake_case      :							str     				=						old_checkpoint[path['''old''']]
                _snake_case      :							Dict     				=						['''proj_attn.weight''']
                _snake_case      :							Union[str, Any]     				=						['''proj_out.weight''', '''proj_in.weight''']
                if any(k in new_path for k in names  ):
                       _snake_case      :							Tuple     				=						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      :							int     				=						weight[:, :, 0]
                else:
                       _snake_case      :							List[str]     				=						weight
  
  
  
  
  def 						A_			(  lowercase_      ,		lowercase_=0  )			->						Optional[int]:
         _snake_case      :							Union[str, Any]     				=						[]
         for old_item in old_list:
                _snake_case      :							List[Any]     				=						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 						A_			(  lowercase_      ,		lowercase_=1  )			->						int:
         if n_shave_prefix_segments >= 0:
                return ".".join(path.split('''.'''  )[n_shave_prefix_segments:]  )
         else:
                return ".".join(path.split('''.'''  )[:n_shave_prefix_segments]  )
  
  
  
  
  def 						A_			(  lowercase_      ,		lowercase_=0  )			->						List[Any]:
         _snake_case      :							Optional[Any]     				=						[]
         for old_item in old_list:
                mapping.append({'''old''': old_item, '''new''': old_item}  )
  
         return mapping
  
  
  
  
  def 						A_			(  lowercase_      ,		lowercase_=0  )			->						str:
         _snake_case      :							Union[str, Any]     				=						[]
         for old_item in old_list:
                _snake_case      :							int     				=						old_item.replace('''in_layers.0'''      ,		'''norm1'''  )
                _snake_case      :							Tuple     				=						new_item.replace('''in_layers.2'''      ,		'''conv1'''  )
  
                _snake_case      :							Optional[int]     				=						new_item.replace('''out_layers.0'''      ,		'''norm2'''  )
                _snake_case      :							Tuple     				=						new_item.replace('''out_layers.3'''      ,		'''conv2'''  )
  
                _snake_case      :							Union[str, Any]     				=						new_item.replace('''emb_layers.1'''      ,		'''time_emb_proj'''  )
                _snake_case      :							Any     				=						new_item.replace('''skip_connection'''      ,		'''conv_shortcut'''  )
  
                _snake_case      :							Optional[Any]     				=						shave_segments(lowercase_      ,		n_shave_prefix_segments=lowercase_  )
  
                if "temopral_conv" not in old_item:
                       mapping.append({'''old''': old_item, '''new''': new_item}  )
  
         return mapping
  
  
  
  
  def 						A_			(  lowercase_      ,		lowercase_      ,		lowercase_=None      ,		lowercase_=False  )			->						Optional[int]:
         _snake_case      :							str     				=						{}
         _snake_case      :							List[str]     				=						list(checkpoint.keys()  )
  
         _snake_case      :							Union[str, Any]     				=						'''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  ) > 100 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      :							Optional[Any]     				=						'''model_ema.''' + ''''''.join(key.split('''.'''  )[1:]  )
                              _snake_case      :							List[str]     				=						checkpoint.pop(lowercase_  )
         else:
                if sum(k.startswith('''model_ema'''  ) for k in keys  ) > 100:
                       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      :							Tuple     				=						checkpoint.pop(lowercase_  )
  
         _snake_case      :							Tuple     				=						{}
  
         _snake_case      :							Tuple     				=						unet_state_dict['''time_embed.0.weight''']
         _snake_case      :							Optional[int]     				=						unet_state_dict['''time_embed.0.bias''']
         _snake_case      :							Any     				=						unet_state_dict['''time_embed.2.weight''']
         _snake_case      :							List[Any]     				=						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      :							int     				=						unet_state_dict['''label_emb.0.0.weight''']
                _snake_case      :							List[Any]     				=						unet_state_dict['''label_emb.0.0.bias''']
                _snake_case      :							Tuple     				=						unet_state_dict['''label_emb.0.2.weight''']
                _snake_case      :							Union[str, Any]     				=						unet_state_dict['''label_emb.0.2.bias''']
         else:
                raise NotImplementedError(f'''Not implemented `class_embed_type`: {config['class_embed_type']}'''  )
  
         _snake_case      :							List[Any]     				=						unet_state_dict['''input_blocks.0.0.weight''']
         _snake_case      :							Tuple     				=						unet_state_dict['''input_blocks.0.0.bias''']
  
         _snake_case      :							Dict     				=						[v for v in unet_state_dict if v.startswith('''input_blocks.0.1'''  )]
         _snake_case      :							Dict     				=						renew_attention_paths(lowercase_  )
         _snake_case      :							str     				=						{'''old''': '''input_blocks.0.1''', '''new''': '''transformer_in'''}
         assign_to_checkpoint(lowercase_      ,		lowercase_      ,		lowercase_      ,		additional_replacements=[meta_path]      ,		config=lowercase_  )
  
         _snake_case      :							str     				=						unet_state_dict['''out.0.weight''']
         _snake_case      :							Any     				=						unet_state_dict['''out.0.bias''']
         _snake_case      :							int     				=						unet_state_dict['''out.2.weight''']
         _snake_case      :							Optional[int]     				=						unet_state_dict['''out.2.bias''']
  
         # Retrieves the keys for the input blocks only
         _snake_case      :							Optional[int]     				=						len({'''.'''.join(layer.split('''.'''  )[:2]  ) for layer in unet_state_dict if '''input_blocks''' in layer}  )
         _snake_case      :							List[str]     				=						{
             layer_id: [key for key in unet_state_dict if f'''input_blocks.{layer_id}''' in key]
             for layer_id in range(lowercase_  )
         }
  
         # Retrieves the keys for the middle blocks only
         _snake_case      :							List[Any]     				=						len({'''.'''.join(layer.split('''.'''  )[:2]  ) for layer in unet_state_dict if '''middle_block''' in layer}  )
         _snake_case      :							Union[str, Any]     				=						{
             layer_id: [key for key in unet_state_dict if f'''middle_block.{layer_id}''' in key]
             for layer_id in range(lowercase_  )
         }
  
         # Retrieves the keys for the output blocks only
         _snake_case      :							Union[str, Any]     				=						len({'''.'''.join(layer.split('''.'''  )[:2]  ) for layer in unet_state_dict if '''output_blocks''' in layer}  )
         _snake_case      :							Tuple     				=						{
             layer_id: [key for key in unet_state_dict if f'''output_blocks.{layer_id}''' in key]
             for layer_id in range(lowercase_  )
         }
  
         for i in range(1      ,		lowercase_  ):
                _snake_case      :							Tuple     				=						(i - 1) // (config['''layers_per_block'''] + 1)
                _snake_case      :							str     				=						(i - 1) % (config['''layers_per_block'''] + 1)
  
                _snake_case      :							int     				=						[
                    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      :							Optional[Any]     				=						[key for key in input_blocks[i] if f'''input_blocks.{i}.1''' in key]
                _snake_case      :							Optional[int]     				=						[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      :							List[str]     				=						unet_state_dict.pop(
                           f'''input_blocks.{i}.op.weight'''  )
                       _snake_case      :							Any     				=						unet_state_dict.pop(
                           f'''input_blocks.{i}.op.bias'''  )
  
                _snake_case      :							str     				=						renew_resnet_paths(lowercase_  )
                _snake_case      :							List[Any]     				=						{'''old''': f'''input_blocks.{i}.0''', '''new''': f'''down_blocks.{block_id}.resnets.{layer_in_block_id}'''}
                assign_to_checkpoint(
                    lowercase_      ,		lowercase_      ,		lowercase_      ,		additional_replacements=[meta_path]      ,		config=lowercase_  )
  
                _snake_case      :							Union[str, Any]     				=						[key for key in resnets if '''temopral_conv''' in key]
                _snake_case      :							int     				=						renew_temp_conv_paths(lowercase_  )
                _snake_case      :							int     				=						{
                    '''old''': f'''input_blocks.{i}.0.temopral_conv''',
                    '''new''': f'''down_blocks.{block_id}.temp_convs.{layer_in_block_id}''',
                }
                assign_to_checkpoint(
                    lowercase_      ,		lowercase_      ,		lowercase_      ,		additional_replacements=[meta_path]      ,		config=lowercase_  )
  
                if len(lowercase_  ):
                       _snake_case      :							Dict     				=						renew_attention_paths(lowercase_  )
                       _snake_case      :							Union[str, Any]     				=						{'''old''': f'''input_blocks.{i}.1''', '''new''': f'''down_blocks.{block_id}.attentions.{layer_in_block_id}'''}
                       assign_to_checkpoint(
                           lowercase_      ,		lowercase_      ,		lowercase_      ,		additional_replacements=[meta_path]      ,		config=lowercase_  )
  
                if len(lowercase_  ):
                       _snake_case      :							Any     				=						renew_attention_paths(lowercase_  )
                       _snake_case      :							Tuple     				=						{
                           '''old''': f'''input_blocks.{i}.2''',
                           '''new''': f'''down_blocks.{block_id}.temp_attentions.{layer_in_block_id}''',
                       }
                       assign_to_checkpoint(
                           lowercase_      ,		lowercase_      ,		lowercase_      ,		additional_replacements=[meta_path]      ,		config=lowercase_  )
  
         _snake_case      :							Dict     				=						middle_blocks[0]
         _snake_case      :							List[str]     				=						[key for key in resnet_a if '''temopral_conv''' in key]
         _snake_case      :							int     				=						middle_blocks[1]
         _snake_case      :							int     				=						middle_blocks[2]
         _snake_case      :							List[str]     				=						middle_blocks[3]
         _snake_case      :							List[Any]     				=						[key for key in resnet_a if '''temopral_conv''' in key]
  
         _snake_case      :							List[str]     				=						renew_resnet_paths(lowercase_  )
         _snake_case      :							Any     				=						{'''old''': '''middle_block.0''', '''new''': '''mid_block.resnets.0'''}
         assign_to_checkpoint(
             lowercase_      ,		lowercase_      ,		lowercase_      ,		config=lowercase_      ,		additional_replacements=[meta_path]  )
  
         _snake_case      :							Optional[Any]     				=						renew_temp_conv_paths(lowercase_  )
         _snake_case      :							Any     				=						{'''old''': '''middle_block.0.temopral_conv''', '''new''': '''mid_block.temp_convs.0'''}
         assign_to_checkpoint(
             lowercase_      ,		lowercase_      ,		lowercase_      ,		config=lowercase_      ,		additional_replacements=[meta_path]  )
  
         _snake_case      :							Tuple     				=						renew_resnet_paths(lowercase_  )
         _snake_case      :							Tuple     				=						{'''old''': '''middle_block.3''', '''new''': '''mid_block.resnets.1'''}
         assign_to_checkpoint(
             lowercase_      ,		lowercase_      ,		lowercase_      ,		config=lowercase_      ,		additional_replacements=[meta_path]  )
  
         _snake_case      :							List[str]     				=						renew_temp_conv_paths(lowercase_  )
         _snake_case      :							Tuple     				=						{'''old''': '''middle_block.3.temopral_conv''', '''new''': '''mid_block.temp_convs.1'''}
         assign_to_checkpoint(
             lowercase_      ,		lowercase_      ,		lowercase_      ,		config=lowercase_      ,		additional_replacements=[meta_path]  )
  
         _snake_case      :							Any     				=						renew_attention_paths(lowercase_  )
         _snake_case      :							str     				=						{'''old''': '''middle_block.1''', '''new''': '''mid_block.attentions.0'''}
         assign_to_checkpoint(
             lowercase_      ,		lowercase_      ,		lowercase_      ,		additional_replacements=[meta_path]      ,		config=lowercase_  )
  
         _snake_case      :							Optional[Any]     				=						renew_attention_paths(lowercase_  )
         _snake_case      :							Any     				=						{'''old''': '''middle_block.2''', '''new''': '''mid_block.temp_attentions.0'''}
         assign_to_checkpoint(
             lowercase_      ,		lowercase_      ,		lowercase_      ,		additional_replacements=[meta_path]      ,		config=lowercase_  )
  
         for i in range(lowercase_  ):
                _snake_case      :							Any     				=						i // (config['''layers_per_block'''] + 1)
                _snake_case      :							Dict     				=						i % (config['''layers_per_block'''] + 1)
                _snake_case      :							Union[str, Any]     				=						[shave_segments(lowercase_      ,		2  ) for name in output_blocks[i]]
                _snake_case      :							Optional[int]     				=						{}
  
                for layer in output_block_layers:
                       _snake_case				,    _snake_case      :							str     				=						layer.split('''.'''  )[0], shave_segments(lowercase_      ,		1  )
                       if layer_id in output_block_list:
                              output_block_list[layer_id].append(lowercase_  )
                       else:
                              _snake_case      :							Any     				=						[layer_name]
  
                if len(lowercase_  ) > 1:
                       _snake_case      :							Optional[int]     				=						[key for key in output_blocks[i] if f'''output_blocks.{i}.0''' in key]
                       _snake_case      :							int     				=						[key for key in output_blocks[i] if f'''output_blocks.{i}.1''' in key]
                       _snake_case      :							Dict     				=						[key for key in output_blocks[i] if f'''output_blocks.{i}.2''' in key]
  
                       _snake_case      :							Tuple     				=						renew_resnet_paths(lowercase_  )
                       _snake_case      :							Optional[Any]     				=						renew_resnet_paths(lowercase_  )
  
                       _snake_case      :							int     				=						{'''old''': f'''output_blocks.{i}.0''', '''new''': f'''up_blocks.{block_id}.resnets.{layer_in_block_id}'''}
                       assign_to_checkpoint(
                           lowercase_      ,		lowercase_      ,		lowercase_      ,		additional_replacements=[meta_path]      ,		config=lowercase_  )
  
                       _snake_case      :							Union[str, Any]     				=						[key for key in resnets if '''temopral_conv''' in key]
                       _snake_case      :							List[Any]     				=						renew_temp_conv_paths(lowercase_  )
                       _snake_case      :							Optional[int]     				=						{
                           '''old''': f'''output_blocks.{i}.0.temopral_conv''',
                           '''new''': f'''up_blocks.{block_id}.temp_convs.{layer_in_block_id}''',
                       }
                       assign_to_checkpoint(
                           lowercase_      ,		lowercase_      ,		lowercase_      ,		additional_replacements=[meta_path]      ,		config=lowercase_  )
  
                       _snake_case      :							Optional[int]     				=						{k: sorted(lowercase_  ) for k, v in output_block_list.items()}
                       if ["conv.bias", "conv.weight"] in output_block_list.values():
                              _snake_case      :							List[Any]     				=						list(output_block_list.values()  ).index(['''conv.bias''', '''conv.weight''']  )
                              _snake_case      :							List[str]     				=						unet_state_dict[
                                  f'''output_blocks.{i}.{index}.conv.weight'''
                              ]
                              _snake_case      :							Any     				=						unet_state_dict[
                                  f'''output_blocks.{i}.{index}.conv.bias'''
                              ]
  
                              # Clear attentions as they have been attributed above.
                              if len(lowercase_  ) == 2:
                                     _snake_case      :							List[str]     				=						[]
  
                       if len(lowercase_  ):
                              _snake_case      :							Union[str, Any]     				=						renew_attention_paths(lowercase_  )
                              _snake_case      :							Tuple     				=						{
                                  '''old''': f'''output_blocks.{i}.1''',
                                  '''new''': f'''up_blocks.{block_id}.attentions.{layer_in_block_id}''',
                              }
                              assign_to_checkpoint(
                                  lowercase_      ,		lowercase_      ,		lowercase_      ,		additional_replacements=[meta_path]      ,		config=lowercase_  )
  
                       if len(lowercase_  ):
                              _snake_case      :							Tuple     				=						renew_attention_paths(lowercase_  )
                              _snake_case      :							Optional[int]     				=						{
                                  '''old''': f'''output_blocks.{i}.2''',
                                  '''new''': f'''up_blocks.{block_id}.temp_attentions.{layer_in_block_id}''',
                              }
                              assign_to_checkpoint(
                                  lowercase_      ,		lowercase_      ,		lowercase_      ,		additional_replacements=[meta_path]      ,		config=lowercase_  )
                else:
                       _snake_case      :							Union[str, Any]     				=						renew_resnet_paths(lowercase_      ,		n_shave_prefix_segments=1  )
                       for path in resnet_0_paths:
                              _snake_case      :							Optional[int]     				=						'''.'''.join(['''output_blocks''', str(lowercase_  ), path['''old''']]  )
                              _snake_case      :							Union[str, Any]     				=						'''.'''.join(['''up_blocks''', str(lowercase_  ), '''resnets''', str(lowercase_  ), path['''new''']]  )
                              _snake_case      :							Tuple     				=						unet_state_dict[old_path]
  
                       _snake_case      :							Tuple     				=						[l for l in output_block_layers if '''temopral_conv''' in l]
                       for path in temopral_conv_paths:
                              _snake_case      :							Union[str, Any]     				=						path.split('''temopral_conv.'''  )[-1]
                              _snake_case      :							Union[str, Any]     				=						'''.'''.join(['''output_blocks''', str(lowercase_  ), str(lowercase_  ), '''temopral_conv''', pruned_path]  )
                              _snake_case      :							int     				=						'''.'''.join(['''up_blocks''', str(lowercase_  ), '''temp_convs''', str(lowercase_  ), pruned_path]  )
                              _snake_case      :							Optional[Any]     				=						unet_state_dict[old_path]
  
         return new_checkpoint
  
  
  if __name__ == "__main__":
     lowerCAmelCase_   =   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.")
     lowerCAmelCase_   =   parser.parse_args()
  
     lowerCAmelCase_   =   torch.load(args.checkpoint_path, map_location="cpu")
     lowerCAmelCase_   =   UNetaDConditionModel()
  
     lowerCAmelCase_   =   convert_ldm_unet_checkpoint(unet_checkpoint, unet.config)
  
     lowerCAmelCase_   =   set(unet.state_dict().keys()) - set(converted_ckpt.keys())
     lowerCAmelCase_   =   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 --
  
  
  
  
  
  
  
style_context: 
  
  
  
  
  
  def 						A_			(  lowercase_  )			->						list:
         if len(lowercase_  ) <= 1:
                return data
         else:
                return [
                    *quick_sort([e for e in data[1:] if e <= data[0]]  ),
                    data[0],
                    *quick_sort([e for e in data[1:] if e > data[0]]  ),
                ]
  
  
  if __name__ == "__main__":
     import doctest
  
     doctest.testmod()
  
  
  
  
  
  
  