




"""simple docstring"""


import inspect
import unittest

from huggingface_hub import hf_hub_download

from transformers import ConvNextConfig, UperNetConfig
from transformers.testing_utils import require_torch, require_torch_multi_gpu, require_vision, slow, torch_device
from transformers.utils import is_torch_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
     import torch

     from transformers import UperNetForSemanticSegmentation
     from transformers.models.upernet.modeling_upernet import UPERNET_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
     from PIL import Image

     from transformers import AutoImageProcessor






class       lowercase:







 '''simple docstring'''





 def __init__(       self:	Dict,   a_:	Union[str, Any],   a_:	Tuple=13,   a_:	Dict=32,   a_:	Optional[Any]=3,   a_:	Optional[Any]=4,   a_:	Optional[int]=[10, 20, 30, 40],   a_:	Any=[2, 2, 3, 2],   a_:	Dict=True,   a_:	Dict=True,   a_:	List[str]=37,   a_:	Dict="gelu",   a_:	List[str]=10,   a_:	Union[str, Any]=0.02,   a_:	Any=["stage2", "stage3", "stage4"],   a_:	Optional[int]=3,   a_:	Tuple=None,   ):




   '''simple docstring'''



   _snake_case							:	Dict			=						parent
   _snake_case							:	Dict			=						batch_size
   _snake_case							:	Optional[Any]			=						image_size
   _snake_case							:	int			=						num_channels
   _snake_case							:	Tuple			=						num_stages
   _snake_case							:	int			=						hidden_sizes
   _snake_case							:	List[str]			=						depths
   _snake_case							:	str			=						is_training
   _snake_case							:	Dict			=						use_labels
   _snake_case							:	List[str]			=						intermediate_size
   _snake_case							:	Optional[int]			=						hidden_act
   _snake_case							:	Any			=						type_sequence_label_size
   _snake_case							:	List[str]			=						initializer_range
   _snake_case							:	Union[str, Any]			=						out_features
   _snake_case							:	Dict			=						num_labels
   _snake_case							:	int			=						scope
   _snake_case							:	Dict			=						num_stages


 def 				UpperCamelCase_     (       self:	int	):




   '''simple docstring'''



   _snake_case							:	Optional[Any]			=						floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]	)

   _snake_case							:	Optional[int]			=						None
   if self.use_labels:
     _snake_case							:	Optional[Any]			=						ids_tensor([self.batch_size],   self.type_sequence_label_size	)

   _snake_case							:	Tuple			=						self.get_config()

   return config, pixel_values, labels


 def 				UpperCamelCase_     (       self:	Tuple	):




   '''simple docstring'''



   return ConvNextConfig(
       num_channels=self.num_channels,   num_stages=self.num_stages,   hidden_sizes=self.hidden_sizes,   depths=self.depths,   is_training=self.is_training,   intermediate_size=self.intermediate_size,   hidden_act=self.hidden_act,   out_features=self.out_features,   )


 def 				UpperCamelCase_     (       self:	List[Any]	):




   '''simple docstring'''



   return UperNetConfig(
       backbone_config=self.get_backbone_config(),   hidden_size=512,   pool_scales=[1, 2, 3, 6],   use_auxiliary_head=a_,   auxiliary_loss_weight=0.4,   auxiliary_in_channels=40,   auxiliary_channels=256,   auxiliary_num_convs=1,   auxiliary_concat_input=a_,   loss_ignore_index=255,   num_labels=self.num_labels,   )


 def 				UpperCamelCase_     (       self:	Tuple,   a_:	List[Any],   a_:	Dict,   a_:	Tuple	):




   '''simple docstring'''



   _snake_case							:	List[Any]			=						UperNetForSemanticSegmentation(config=a_	)
   model.to(a_	)
   model.eval()
   _snake_case							:	Tuple			=						model(a_	)
   self.parent.assertEqual(
       result.logits.shape,   (self.batch_size, self.num_labels, self.image_size, self.image_size)	)


 def 				UpperCamelCase_     (       self:	Optional[Any]	):




   '''simple docstring'''



   _snake_case							:	Any			=						self.prepare_config_and_inputs()
   (
       (
       _snake_case
   )	,      (
       _snake_case
   )	,      (
       _snake_case
   )	,      
   )							:	List[Any]			=						config_and_inputs
   _snake_case							:	Any			=						{"""pixel_values""": pixel_values}
   return config, inputs_dict







@require_torch
class       lowercase(  __a ,    __a ,    unittest.TestCase				):







 '''simple docstring'''





 lowercase__				      =      (UperNetForSemanticSegmentation,) if is_torch_available() else ()
 lowercase__				      =      {"image-segmentation": UperNetForSemanticSegmentation} if is_torch_available() else {}
 lowercase__				      =      False
 lowercase__				      =      False
 lowercase__				      =      False
 lowercase__				      =      False
 lowercase__				      =      False
 lowercase__				      =      False


 def 				UpperCamelCase_     (       self:	Optional[Any]	):




   '''simple docstring'''



   _snake_case							:	List[str]			=						UperNetModelTester(self	)
   _snake_case							:	Dict			=						ConfigTester(self,   config_class=a_,   has_text_modality=a_,   hidden_size=37	)


 def 				UpperCamelCase_     (       self:	Any	):




   '''simple docstring'''



   self.create_and_test_config_common_properties()
   self.config_tester.create_and_test_config_to_json_string()
   self.config_tester.create_and_test_config_to_json_file()
   self.config_tester.create_and_test_config_from_and_save_pretrained()
   self.config_tester.create_and_test_config_with_num_labels()
   self.config_tester.check_config_can_be_init_without_params()
   self.config_tester.check_config_arguments_init()


 def 				UpperCamelCase_     (       self:	List[Any]	):




   '''simple docstring'''



   return


 def 				UpperCamelCase_     (       self:	Union[str, Any]	):




   '''simple docstring'''



   _snake_case	,      _snake_case							:	Any			=						self.model_tester.prepare_config_and_inputs_for_common()

   for model_class in self.all_model_classes:
     _snake_case							:	Dict			=						model_class(a_	)
     _snake_case							:	Optional[int]			=						inspect.signature(model.forward	)
     # signature.parameters is an OrderedDict => so arg_names order is deterministic
     _snake_case							:	Tuple			=						[*signature.parameters.keys()]

     _snake_case							:	Any			=						["""pixel_values"""]
     self.assertListEqual(arg_names[:1],   a_	)


 def 				UpperCamelCase_     (       self:	Any	):




   '''simple docstring'''



   _snake_case							:	Union[str, Any]			=						self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_semantic_segmentation(*a_	)


 @unittest.skip(reason="""UperNet does not use inputs_embeds"""	)
 def 				UpperCamelCase_     (       self:	Any	):




   '''simple docstring'''



   pass


 @unittest.skip(reason="""UperNet does not support input and output embeddings"""	)
 def 				UpperCamelCase_     (       self:	Dict	):




   '''simple docstring'''



   pass


 @unittest.skip(reason="""UperNet does not have a base model"""	)
 def 				UpperCamelCase_     (       self:	Optional[Any]	):




   '''simple docstring'''



   pass


 @unittest.skip(reason="""UperNet does not have a base model"""	)
 def 				UpperCamelCase_     (       self:	int	):




   '''simple docstring'''



   pass


 @require_torch_multi_gpu
 @unittest.skip(reason="""UperNet has some layers using `add_module` which doesn't work well with `nn.DataParallel`"""	)
 def 				UpperCamelCase_     (       self:	List[str]	):




   '''simple docstring'''



   pass


 @unittest.skip("""Will be fixed soon by reducing the size of the model used for common tests."""	)
 def 				UpperCamelCase_     (       self:	Dict	):




   '''simple docstring'''



   pass


 def 				UpperCamelCase_     (       self:	str	):




   '''simple docstring'''



   def check_hidden_states_output(a_:	Dict,   a_:	List[str],   a_:	Optional[int]	):
     _snake_case							:	Optional[Any]			=						model_class(a_	)
     model.to(a_	)
     model.eval()

     with torch.no_grad():
       _snake_case							:	Any			=						model(**self._prepare_for_class(a_,   a_	)	)

     _snake_case							:	Any			=						outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

     _snake_case							:	List[str]			=						self.model_tester.num_stages
     self.assertEqual(len(a_	),   expected_num_stages + 1	)

     # ConvNext's feature maps are of shape (batch_size, num_channels, height, width)
     self.assertListEqual(
         list(hidden_states[0].shape[-2:]	),   [self.model_tester.image_size // 4, self.model_tester.image_size // 4],   )

   _snake_case	,      _snake_case							:	str			=						self.model_tester.prepare_config_and_inputs_for_common()

   for model_class in self.all_model_classes:
     _snake_case							:	int			=						True
     check_hidden_states_output(a_,   a_,   a_	)

     # check that output_hidden_states also work using config
     del inputs_dict["output_hidden_states"]
     _snake_case							:	Optional[int]			=						True

     check_hidden_states_output(a_,   a_,   a_	)


 def 				UpperCamelCase_     (       self:	str	):




   '''simple docstring'''



   _snake_case	,      _snake_case							:	int			=						self.model_tester.prepare_config_and_inputs_for_common()

   _snake_case							:	Tuple			=						_config_zero_init(a_	)
   _snake_case							:	Dict			=						_config_zero_init(configs_no_init.backbone_config	)
   for model_class in self.all_model_classes:
     _snake_case							:	Optional[int]			=						model_class(config=a_	)
     for name, param in model.named_parameters():
       if param.requires_grad:
         self.assertIn(
             ((param.data.mean() * 1E9).round() / 1E9).item(),   [0.0, 1.0],   msg=f"Parameter {name} of model {model_class} seems not properly initialized",   )


 @unittest.skip(reason="""UperNet does not have tied weights"""	)
 def 				UpperCamelCase_     (       self:	Union[str, Any]	):




   '''simple docstring'''



   pass


 @slow
 def 				UpperCamelCase_     (       self:	List[str]	):




   '''simple docstring'''



   for model_name in UPERNET_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
     _snake_case							:	int			=						UperNetForSemanticSegmentation.from_pretrained(a_	)
     self.assertIsNotNone(a_	)







def       UpperCAmelCase__ ():
  """simple docstring"""

  _snake_case							:	Union[str, Any]			=						hf_hub_download(
      repo_id="""hf-internal-testing/fixtures_ade20k"""	,     repo_type="""dataset"""	,     filename="""ADE_val_00000001.jpg"""					)
  _snake_case							:	List[Any]			=						Image.open(snake_case__					).convert("""RGB"""					)
  return image






@require_torch
@require_vision
@slow
class       lowercase(  unittest.TestCase				):







 '''simple docstring'''





 def 				UpperCamelCase_     (       self:	Tuple	):




   '''simple docstring'''



   _snake_case							:	str			=						AutoImageProcessor.from_pretrained("""openmmlab/upernet-swin-tiny"""	)
   _snake_case							:	Any			=						UperNetForSemanticSegmentation.from_pretrained("""openmmlab/upernet-swin-tiny"""	).to(a_	)

   _snake_case							:	Dict			=						prepare_img()
   _snake_case							:	str			=						processor(images=a_,   return_tensors="""pt"""	).to(a_	)

   with torch.no_grad():
     _snake_case							:	Tuple			=						model(**a_	)

   _snake_case							:	Tuple			=						torch.Size((1, model.config.num_labels, 512, 512)	)
   self.assertEqual(outputs.logits.shape,   a_	)

   _snake_case							:	int			=						torch.tensor(
       [[-7.5_958, -7.5_958, -7.4_302], [-7.5_958, -7.5_958, -7.4_302], [-7.4_797, -7.4_797, -7.3_068]]	).to(a_	)
   self.assertTrue(torch.allclose(outputs.logits[0, 0, :3, :3],   a_,   atol=1E-4	)	)


 def 				UpperCamelCase_     (       self:	int	):




   '''simple docstring'''



   _snake_case							:	Optional[Any]			=						AutoImageProcessor.from_pretrained("""openmmlab/upernet-convnext-tiny"""	)
   _snake_case							:	Optional[int]			=						UperNetForSemanticSegmentation.from_pretrained("""openmmlab/upernet-convnext-tiny"""	).to(a_	)

   _snake_case							:	List[str]			=						prepare_img()
   _snake_case							:	Tuple			=						processor(images=a_,   return_tensors="""pt"""	).to(a_	)

   with torch.no_grad():
     _snake_case							:	Optional[Any]			=						model(**a_	)

   _snake_case							:	Union[str, Any]			=						torch.Size((1, model.config.num_labels, 512, 512)	)
   self.assertEqual(outputs.logits.shape,   a_	)

   _snake_case							:	Optional[Any]			=						torch.tensor(
       [[-8.8_110, -8.8_110, -8.6_521], [-8.8_110, -8.8_110, -8.6_521], [-8.7_746, -8.7_746, -8.6_130]]	).to(a_	)
   self.assertTrue(torch.allclose(outputs.logits[0, 0, :3, :3],   a_,   atol=1E-4	)	)








"""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)


