code: 
  
  
  
  
  
  import argparse
  import json
  
  import torch
  
  from diffusers import ScoreSdeVePipeline, ScoreSdeVeScheduler, UNetaDModel
  
  
  
  
  
  
  
  def       lowerCAmelCase__(	lowercase       :    Tuple   ,     lowercase       :    Union[str, Any]				) ->    Union[str, Any]:
    __snake_case   :				List[str]        =       UNetaDModel(**lowercase				)
    __snake_case   :				int        =       checkpoint["all_modules.0.W"].data
    __snake_case   :				List[Any]        =       checkpoint["all_modules.0.W"].data
    __snake_case   :				List[str]        =       checkpoint["all_modules.1.weight"].data
    __snake_case   :				List[Any]        =       checkpoint["all_modules.1.bias"].data
  
    __snake_case   :				int        =       checkpoint["all_modules.2.weight"].data
    __snake_case   :				Any        =       checkpoint["all_modules.2.bias"].data
  
    __snake_case   :				Optional[Any]        =       checkpoint["all_modules.3.weight"].data
    __snake_case   :				List[Any]        =       checkpoint["all_modules.3.bias"].data
  
    __snake_case   :				Dict        =       checkpoint[list(checkpoint.keys()				)[-4]].data
    __snake_case   :				int        =       checkpoint[list(checkpoint.keys()				)[-3]].data
    __snake_case   :				List[Any]        =       checkpoint[list(checkpoint.keys()				)[-2]].data
    __snake_case   :				Optional[int]        =       checkpoint[list(checkpoint.keys()				)[-1]].data
  
    __snake_case   :				int        =       4
  
    def set_attention_weights(lowercase       :    List[str]   ,     lowercase       :    Dict   ,     lowercase       :    Dict				):
      __snake_case   :				Any        =       old_checkpoint[f"""all_modules.{index}.NIN_0.W"""].data.T
      __snake_case   :				Dict        =       old_checkpoint[f"""all_modules.{index}.NIN_1.W"""].data.T
      __snake_case   :				str        =       old_checkpoint[f"""all_modules.{index}.NIN_2.W"""].data.T
  
      __snake_case   :				str        =       old_checkpoint[f"""all_modules.{index}.NIN_0.b"""].data
      __snake_case   :				Optional[Any]        =       old_checkpoint[f"""all_modules.{index}.NIN_1.b"""].data
      __snake_case   :				Dict        =       old_checkpoint[f"""all_modules.{index}.NIN_2.b"""].data
  
      __snake_case   :				List[Any]        =       old_checkpoint[f"""all_modules.{index}.NIN_3.W"""].data.T
      __snake_case   :				Any        =       old_checkpoint[f"""all_modules.{index}.NIN_3.b"""].data
  
      __snake_case   :				Tuple        =       old_checkpoint[f"""all_modules.{index}.GroupNorm_0.weight"""].data
      __snake_case   :				str        =       old_checkpoint[f"""all_modules.{index}.GroupNorm_0.bias"""].data
  
    def set_resnet_weights(lowercase       :    Optional[Any]   ,     lowercase       :    Optional[Any]   ,     lowercase       :    int				):
      __snake_case   :				Optional[int]        =       old_checkpoint[f"""all_modules.{index}.Conv_0.weight"""].data
      __snake_case   :				Optional[Any]        =       old_checkpoint[f"""all_modules.{index}.Conv_0.bias"""].data
      __snake_case   :				List[str]        =       old_checkpoint[f"""all_modules.{index}.GroupNorm_0.weight"""].data
      __snake_case   :				str        =       old_checkpoint[f"""all_modules.{index}.GroupNorm_0.bias"""].data
  
      __snake_case   :				Union[str, Any]        =       old_checkpoint[f"""all_modules.{index}.Conv_1.weight"""].data
      __snake_case   :				Optional[Any]        =       old_checkpoint[f"""all_modules.{index}.Conv_1.bias"""].data
      __snake_case   :				Dict        =       old_checkpoint[f"""all_modules.{index}.GroupNorm_1.weight"""].data
      __snake_case   :				Any        =       old_checkpoint[f"""all_modules.{index}.GroupNorm_1.bias"""].data
  
      __snake_case   :				Optional[Any]        =       old_checkpoint[f"""all_modules.{index}.Dense_0.weight"""].data
      __snake_case   :				List[Any]        =       old_checkpoint[f"""all_modules.{index}.Dense_0.bias"""].data
  
      if new_layer.in_channels != new_layer.out_channels or new_layer.up or new_layer.down:
        __snake_case   :				Optional[Any]        =       old_checkpoint[f"""all_modules.{index}.Conv_2.weight"""].data
        __snake_case   :				Union[str, Any]        =       old_checkpoint[f"""all_modules.{index}.Conv_2.bias"""].data
  
    for i, block in enumerate(new_model_architecture.downsample_blocks				):
      __snake_case   :				Tuple        =       hasattr(lowercase   ,     "attentions"				)
      for j in range(len(block.resnets				)				):
        set_resnet_weights(block.resnets[j]   ,     lowercase   ,     lowercase				)
        module_index += 1
        if has_attentions:
          set_attention_weights(block.attentions[j]   ,     lowercase   ,     lowercase				)
          module_index += 1
  
      if hasattr(lowercase   ,     "downsamplers"				) and block.downsamplers is not None:
        set_resnet_weights(block.resnet_down   ,     lowercase   ,     lowercase				)
        module_index += 1
        __snake_case   :				List[str]        =       checkpoint[f"""all_modules.{module_index}.Conv_0.weight"""].data
        __snake_case   :				Optional[int]        =       checkpoint[f"""all_modules.{module_index}.Conv_0.bias"""].data
        module_index += 1
  
    set_resnet_weights(new_model_architecture.mid_block.resnets[0]   ,     lowercase   ,     lowercase				)
    module_index += 1
    set_attention_weights(new_model_architecture.mid_block.attentions[0]   ,     lowercase   ,     lowercase				)
    module_index += 1
    set_resnet_weights(new_model_architecture.mid_block.resnets[1]   ,     lowercase   ,     lowercase				)
    module_index += 1
  
    for i, block in enumerate(new_model_architecture.up_blocks				):
      __snake_case   :				Any        =       hasattr(lowercase   ,     "attentions"				)
      for j in range(len(block.resnets				)				):
        set_resnet_weights(block.resnets[j]   ,     lowercase   ,     lowercase				)
        module_index += 1
      if has_attentions:
        set_attention_weights(
            block.attentions[0]   ,     lowercase   ,     lowercase				)  # why can there only be a single attention layer for up?
        module_index += 1
  
      if hasattr(lowercase   ,     "resnet_up"				) and block.resnet_up is not None:
        __snake_case   :				Optional[Any]        =       checkpoint[f"""all_modules.{module_index}.weight"""].data
        __snake_case   :				Any        =       checkpoint[f"""all_modules.{module_index}.bias"""].data
        module_index += 1
        __snake_case   :				Any        =       checkpoint[f"""all_modules.{module_index}.weight"""].data
        __snake_case   :				List[str]        =       checkpoint[f"""all_modules.{module_index}.bias"""].data
        module_index += 1
        set_resnet_weights(block.resnet_up   ,     lowercase   ,     lowercase				)
        module_index += 1
  
    __snake_case   :				Any        =       checkpoint[f"""all_modules.{module_index}.weight"""].data
    __snake_case   :				Any        =       checkpoint[f"""all_modules.{module_index}.bias"""].data
    module_index += 1
    __snake_case   :				Tuple        =       checkpoint[f"""all_modules.{module_index}.weight"""].data
    __snake_case   :				str        =       checkpoint[f"""all_modules.{module_index}.bias"""].data
  
    return new_model_architecture.state_dict()
  
  
  if __name__ == "__main__":
       _UpperCamelCase       = argparse.ArgumentParser()
  
       parser.add_argument(
           '''--checkpoint_path''',
           default='''/Users/arthurzucker/Work/diffusers/ArthurZ/diffusion_pytorch_model.bin''',
           type=str,
           required=False,
           help='''Path to the checkpoint to convert.''',
       )
  
       parser.add_argument(
           '''--config_file''',
           default='''/Users/arthurzucker/Work/diffusers/ArthurZ/config.json''',
           type=str,
           required=False,
           help='''The config json file corresponding to the architecture.''',
       )
  
       parser.add_argument(
           '''--dump_path''',
           default='''/Users/arthurzucker/Work/diffusers/ArthurZ/diffusion_model_new.pt''',
           type=str,
           required=False,
           help='''Path to the output model.''',
       )
  
       _UpperCamelCase       = parser.parse_args()
  
       _UpperCamelCase       = torch.load(args.checkpoint_path, map_location='''cpu''')
  
       with open(args.config_file) as f:
            _UpperCamelCase       = json.loads(f.read())
  
       _UpperCamelCase       = convert_ncsnpp_checkpoint(
           checkpoint,
           config,
       )
  
       if "sde" in config:
            del config["sde"]
  
       _UpperCamelCase       = UNetaDModel(**config)
       model.load_state_dict(converted_checkpoint)
  
       try:
            _UpperCamelCase       = ScoreSdeVeScheduler.from_config('''/'''.join(args.checkpoint_path.split('''/''')[:-1]))
  
            _UpperCamelCase       = ScoreSdeVePipeline(unet=model, scheduler=scheduler)
            pipe.save_pretrained(args.dump_path)
       except:  # noqa: E722
            model.save_pretrained(args.dump_path)
  
  
style_context: 
  
  
  
  
  
  from typing import TYPE_CHECKING
  
  from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available
  
  
  _UpperCamelCase       = {'''configuration_nat''': ['''NAT_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''NatConfig''']}
  
  
  try:
       if not is_torch_available():
            raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
       pass
  else:
       _UpperCamelCase       = [
           '''NAT_PRETRAINED_MODEL_ARCHIVE_LIST''',
           '''NatForImageClassification''',
           '''NatModel''',
           '''NatPreTrainedModel''',
           '''NatBackbone''',
       ]
  
  if TYPE_CHECKING:
       from .configuration_nat import NAT_PRETRAINED_CONFIG_ARCHIVE_MAP, NatConfig
  
       try:
            if not is_torch_available():
                 raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
            pass
       else:
            from .modeling_nat import (
                NAT_PRETRAINED_MODEL_ARCHIVE_LIST,
                NatBackbone,
                NatForImageClassification,
                NatModel,
                NatPreTrainedModel,
            )
  
  else:
       import sys
  
       _UpperCamelCase       = _LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)
  
  