



"""simple docstring"""




import argparse
import json

import torch

from diffusers import DDPMScheduler, LDMPipeline, UNetaDModel, VQModel

def 		__snake_case		( SCREAMING_SNAKE_CASE__      :		str		,    SCREAMING_SNAKE_CASE__      :		Any=1							)    ->  int:





     '''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 		__snake_case		( SCREAMING_SNAKE_CASE__      :		Any		,    SCREAMING_SNAKE_CASE__      :		Dict=0							)    ->  int:





     '''simple docstring'''

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

          _UpperCAmelCase     :      Tuple														=       new_item.replace("out_layers.0"		,    "norm2"							)
          _UpperCAmelCase     :      Optional[Any]														=       new_item.replace("out_layers.3"		,    "conv2"							)

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

          _UpperCAmelCase     :      List[Any]														=       shave_segments(SCREAMING_SNAKE_CASE__		,    n_shave_prefix_segments=SCREAMING_SNAKE_CASE__							)

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

     return mapping

def 		__snake_case		( SCREAMING_SNAKE_CASE__      :		List[Any]		,    SCREAMING_SNAKE_CASE__      :		int=0							)    ->  int:





     '''simple docstring'''

     _UpperCAmelCase     :      str														=       []
     for old_item in old_list:
          _UpperCAmelCase     :      Dict														=       old_item

          _UpperCAmelCase     :      Optional[int]														=       new_item.replace("norm.weight"		,    "group_norm.weight"							)
          _UpperCAmelCase     :      Optional[Any]														=       new_item.replace("norm.bias"		,    "group_norm.bias"							)

          _UpperCAmelCase     :      str														=       new_item.replace("proj_out.weight"		,    "proj_attn.weight"							)
          _UpperCAmelCase     :      List[str]														=       new_item.replace("proj_out.bias"		,    "proj_attn.bias"							)

          _UpperCAmelCase     :      str														=       shave_segments(SCREAMING_SNAKE_CASE__		,    n_shave_prefix_segments=SCREAMING_SNAKE_CASE__							)

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

     return mapping

def 		__snake_case		( SCREAMING_SNAKE_CASE__      :		Optional[Any]		,    SCREAMING_SNAKE_CASE__      :		Optional[Any]		,    SCREAMING_SNAKE_CASE__      :		Tuple		,    SCREAMING_SNAKE_CASE__      :		Optional[int]=None		,    SCREAMING_SNAKE_CASE__      :		Tuple=None		,    SCREAMING_SNAKE_CASE__      :		Dict=None							)    ->  Dict:





     '''simple docstring'''

     assert isinstance(SCREAMING_SNAKE_CASE__		,    SCREAMING_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():
               _UpperCAmelCase     :      List[Any]														=       old_checkpoint[path]
               _UpperCAmelCase     :      int														=       old_tensor.shape[0] // 3

               _UpperCAmelCase     :      Tuple														=       (-1, channels) if len(old_tensor.shape							) == 3 else (-1)

               _UpperCAmelCase     :      Dict														=       old_tensor.shape[0] // config["num_head_channels"] // 3

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

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

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

          if additional_replacements is not None:
               for replacement in additional_replacements:
                    _UpperCAmelCase     :      List[str]														=       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:
               _UpperCAmelCase     :      Optional[int]														=       old_checkpoint[path["old"]][:, :, 0]
          else:
               _UpperCAmelCase     :      int														=       old_checkpoint[path["old"]]

def 		__snake_case		( SCREAMING_SNAKE_CASE__      :		Any		,    SCREAMING_SNAKE_CASE__      :		Any							)    ->  Optional[Any]:





     '''simple docstring'''

     _UpperCAmelCase     :      Optional[Any]														=       {}

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

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

     _UpperCAmelCase     :      Any														=       checkpoint["out.0.weight"]
     _UpperCAmelCase     :      Optional[int]														=       checkpoint["out.0.bias"]
     _UpperCAmelCase     :      Tuple														=       checkpoint["out.2.weight"]
     _UpperCAmelCase     :      Optional[int]														=       checkpoint["out.2.bias"]

     # Retrieves the keys for the input blocks only
     _UpperCAmelCase     :      List[Any]														=       len({".".join(layer.split("."							)[:2]							) for layer in checkpoint if "input_blocks" in layer}							)
     _UpperCAmelCase     :      Optional[Any]														=       {
         layer_id: [key for key in checkpoint if f'input_blocks.{layer_id}' in key]
         for layer_id in range(SCREAMING_SNAKE_CASE__							)
     }

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

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

     for i in range(1		,    SCREAMING_SNAKE_CASE__							):
          _UpperCAmelCase     :      Any														=       (i - 1) // (config["num_res_blocks"] + 1)
          _UpperCAmelCase     :      Tuple														=       (i - 1) % (config["num_res_blocks"] + 1)

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

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

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

     _UpperCAmelCase     :      List[Any]														=       middle_blocks[0]
     _UpperCAmelCase     :      List[str]														=       middle_blocks[1]
     _UpperCAmelCase     :      List[str]														=       middle_blocks[2]

     _UpperCAmelCase     :      str														=       renew_resnet_paths(SCREAMING_SNAKE_CASE__							)
     assign_to_checkpoint(SCREAMING_SNAKE_CASE__		,    SCREAMING_SNAKE_CASE__		,    SCREAMING_SNAKE_CASE__		,    config=SCREAMING_SNAKE_CASE__							)

     _UpperCAmelCase     :      Tuple														=       renew_resnet_paths(SCREAMING_SNAKE_CASE__							)
     assign_to_checkpoint(SCREAMING_SNAKE_CASE__		,    SCREAMING_SNAKE_CASE__		,    SCREAMING_SNAKE_CASE__		,    config=SCREAMING_SNAKE_CASE__							)

     _UpperCAmelCase     :      List[str]														=       renew_attention_paths(SCREAMING_SNAKE_CASE__							)
     _UpperCAmelCase     :      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(
         SCREAMING_SNAKE_CASE__		,    SCREAMING_SNAKE_CASE__		,    SCREAMING_SNAKE_CASE__		,    attention_paths_to_split=SCREAMING_SNAKE_CASE__		,    config=SCREAMING_SNAKE_CASE__							)

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

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

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

               _UpperCAmelCase     :      List[str]														=       renew_resnet_paths(SCREAMING_SNAKE_CASE__							)
               _UpperCAmelCase     :      Optional[Any]														=       renew_resnet_paths(SCREAMING_SNAKE_CASE__							)

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

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

                    # Clear attentions as they have been attributed above.
                    if len(SCREAMING_SNAKE_CASE__							) == 2:
                         _UpperCAmelCase     :      Optional[int]														=       []

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

                    _UpperCAmelCase     :      int														=       checkpoint[old_path]

     return new_checkpoint


if __name__ == "__main__":
 _lowerCAmelCase							:   List[str]				=				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.")

 _lowerCAmelCase							:   Tuple				=				parser.parse_args()

 _lowerCAmelCase							:   Optional[int]				=				torch.load(args.checkpoint_path)

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

 _lowerCAmelCase							:   Dict				=				convert_ldm_checkpoint(checkpoint, config)

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

 _lowerCAmelCase							:   Optional[int]				=				UNetaDModel(**config)
 model.load_state_dict(converted_checkpoint)

 try:
  _lowerCAmelCase							:   Tuple				=				DDPMScheduler.from_config("/".join(args.checkpoint_path.split("/")[:-1]))
  _lowerCAmelCase							:   Tuple				=				VQModel.from_pretrained("/".join(args.checkpoint_path.split("/")[:-1]))

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





"""simple docstring"""




import random

import torch
from huggingface_hub import HfApi

from diffusers import UNetaDModel


_lowerCAmelCase							:   Any				=				HfApi()

_lowerCAmelCase							:   int				=				{}
# fmt: off
_lowerCAmelCase							:   List[Any]				=				torch.tensor([
    -0.7_5_1_5, -1.6_8_8_3, 0.2_4_2_0, 0.0_3_0_0, 0.6_3_4_7, 1.3_4_3_3, -1.1_7_4_3, -3.7_4_6_7,
    1.2_3_4_2, -2.2_4_8_5, 0.4_6_3_6, 0.8_0_7_6, -0.7_9_9_1, 0.3_9_6_9, 0.8_4_9_8, 0.9_1_8_9,
    -1.8_8_8_7, -3.3_5_2_2, 0.7_6_3_9, 0.2_0_4_0, 0.6_2_7_1, -2.7_1_4_8, -1.6_3_1_6, 3.0_8_3_9,
    0.3_1_8_6, 0.2_7_2_1, -0.9_7_5_9, -1.2_4_6_1, 2.6_2_5_7, 1.3_5_5_7
])
_lowerCAmelCase							:   Any				=				torch.tensor([
    -2.3_6_3_9, -2.5_3_4_4, 0.0_0_5_4, -0.6_6_7_4, 1.5_9_9_0, 1.0_1_5_8, 0.3_1_2_4, -2.1_4_3_6,
    1.8_7_9_5, -2.5_4_2_9, -0.1_5_6_6, -0.3_9_7_3, 1.2_4_9_0, 2.6_4_4_7, 1.2_2_8_3, -0.5_2_0_8,
    -2.8_1_5_4, -3.5_1_1_9, 2.3_8_3_8, 1.2_0_3_3, 1.7_2_0_1, -2.1_2_5_6, -1.4_5_7_6, 2.7_9_4_8,
    2.4_2_0_4, -0.9_7_5_2, -1.2_5_4_6, 0.8_0_2_7, 3.2_7_5_8, 3.1_3_6_5
])
_lowerCAmelCase							:   List[str]				=				torch.tensor([
    -0.6_5_3_1, -0.6_8_9_1, -0.3_1_7_2, -0.5_3_7_5, -0.9_1_4_0, -0.5_3_6_7, -0.1_1_7_5, -0.7_8_6_9,
    -0.3_8_0_8, -0.4_5_1_3, -0.2_0_9_8, -0.0_0_8_3, 0.3_1_8_3, 0.5_1_4_0, 0.2_2_4_7, -0.1_3_0_4,
    -0.1_3_0_2, -0.2_8_0_2, -0.2_0_8_4, -0.2_0_2_5, -0.4_9_6_7, -0.4_8_7_3, -0.0_8_6_1, 0.6_9_2_5,
    0.0_2_5_0, 0.1_2_9_0, -0.1_5_4_3, 0.6_3_1_6, 1.0_4_6_0, 1.4_9_4_3
])
_lowerCAmelCase							:   Optional[int]				=				torch.tensor([
    0.0_9_1_1, 0.1_1_0_7, 0.0_1_8_2, 0.0_4_3_5, -0.0_8_0_5, -0.0_6_0_8, 0.0_3_8_1, 0.2_1_7_2,
    -0.0_2_8_0, 0.1_3_2_7, -0.0_2_9_9, -0.0_2_5_5, -0.0_0_5_0, -0.1_1_7_0, -0.1_0_4_6, 0.0_3_0_9,
    0.1_3_6_7, 0.1_7_2_8, -0.0_5_3_3, -0.0_7_4_8, -0.0_5_3_4, 0.1_6_2_4, 0.0_3_8_4, -0.1_8_0_5,
    -0.0_7_0_7, 0.0_6_4_2, 0.0_2_2_0, -0.0_1_3_4, -0.1_3_3_3, -0.1_5_0_5
])
_lowerCAmelCase							:   List[Any]				=				torch.tensor([
    0.1_3_2_1, 0.1_3_3_7, 0.0_4_4_0, 0.0_6_2_2, -0.0_5_9_1, -0.0_3_7_0, 0.0_5_0_3, 0.2_1_3_3,
    -0.0_1_7_7, 0.1_4_1_5, -0.0_1_1_6, -0.0_1_1_2, 0.0_0_4_4, -0.0_9_8_0, -0.0_7_8_9, 0.0_3_9_5,
    0.1_5_0_2, 0.1_7_8_5, -0.0_4_8_8, -0.0_5_1_4, -0.0_4_0_4, 0.1_5_3_9, 0.0_4_5_4, -0.1_5_5_9,
    -0.0_6_6_5, 0.0_6_5_9, 0.0_3_8_3, -0.0_0_0_5, -0.1_2_6_6, -0.1_3_8_6
])
_lowerCAmelCase							:   Optional[int]				=				torch.tensor([
    0.1_1_5_4, 0.1_2_1_8, 0.0_3_0_7, 0.0_5_2_6, -0.0_7_1_1, -0.0_5_4_1, 0.0_3_6_6, 0.2_0_7_8,
    -0.0_2_6_7, 0.1_3_1_7, -0.0_2_2_6, -0.0_1_9_3, -0.0_0_1_4, -0.1_0_5_5, -0.0_9_0_2, 0.0_3_3_0,
    0.1_3_9_1, 0.1_7_0_9, -0.0_5_6_2, -0.0_6_9_3, -0.0_5_6_0, 0.1_4_8_2, 0.0_3_8_1, -0.1_6_8_3,
    -0.0_6_8_1, 0.0_6_6_1, 0.0_3_3_1, -0.0_0_4_6, -0.1_2_6_8, -0.1_4_3_1
])
_lowerCAmelCase							:   Optional[Any]				=				torch.tensor([
    0.1_1_9_2, 0.1_2_4_0, 0.0_4_1_4, 0.0_6_0_6, -0.0_5_5_7, -0.0_4_1_2, 0.0_4_3_0, 0.2_0_4_2,
    -0.0_2_0_0, 0.1_3_8_5, -0.0_1_1_5, -0.0_1_3_2, 0.0_0_1_7, -0.0_9_6_5, -0.0_8_0_2, 0.0_3_9_8,
    0.1_4_3_3, 0.1_7_4_7, -0.0_4_5_8, -0.0_5_3_3, -0.0_4_0_7, 0.1_5_4_5, 0.0_4_1_9, -0.1_5_7_4,
    -0.0_6_4_5, 0.0_6_2_6, 0.0_3_4_1, -0.0_0_1_0, -0.1_1_9_9, -0.1_3_9_0
])
_lowerCAmelCase							:   List[str]				=				torch.tensor([
    0.1_0_7_5, 0.1_0_7_4, 0.0_2_0_5, 0.0_4_3_1, -0.0_7_7_4, -0.0_6_0_7, 0.0_2_9_8, 0.2_0_4_2,
    -0.0_3_2_0, 0.1_2_6_7, -0.0_2_8_1, -0.0_2_5_0, -0.0_0_6_4, -0.1_0_9_1, -0.0_9_4_6, 0.0_2_9_0,
    0.1_3_2_8, 0.1_6_5_0, -0.0_5_8_0, -0.0_7_3_8, -0.0_5_8_6, 0.1_4_4_0, 0.0_3_3_7, -0.1_7_4_6,
    -0.0_7_1_2, 0.0_6_0_5, 0.0_2_5_0, -0.0_0_9_9, -0.1_3_1_6, -0.1_4_7_3
])
_lowerCAmelCase							:   Any				=				torch.tensor([
    -1.4_5_7_2, -2.0_4_8_1, -0.0_4_1_4, -0.6_0_0_5, 1.4_1_3_6, 0.5_8_4_8, 0.4_0_2_8, -2.7_3_3_0,
    1.2_2_1_2, -2.1_2_2_8, 0.2_1_5_5, 0.4_0_3_9, 0.7_6_6_2, 2.0_5_3_5, 0.7_4_7_7, -0.3_2_4_3,
    -2.1_7_5_8, -2.7_6_4_8, 1.6_9_4_7, 0.7_0_2_6, 1.2_3_3_8, -1.6_0_7_8, -0.8_6_8_2, 2.2_8_1_0,
    1.8_5_7_4, -0.5_7_1_8, -0.5_5_8_6, -0.0_1_8_6, 2.3_4_1_5, 2.1_2_5_1])
_lowerCAmelCase							:   int				=				torch.tensor([
    -1.3_6_9_0, -1.9_7_2_0, -0.4_0_9_0, -0.6_9_6_6, 1.4_6_6_0, 0.9_9_3_8, -0.1_3_8_5, -2.7_3_2_4,
    0.7_7_3_6, -1.8_9_1_7, 0.2_9_2_3, 0.4_2_9_3, 0.1_6_9_3, 1.4_1_1_2, 1.1_8_8_7, -0.3_1_8_1,
    -2.2_1_6_0, -2.6_3_8_1, 1.3_1_7_0, 0.8_1_6_3, 0.9_2_4_0, -1.6_5_4_4, -0.6_0_9_9, 2.5_2_5_9,
    1.6_4_3_0, -0.9_0_9_0, -0.9_3_9_2, -0.0_1_2_6, 2.4_2_6_8, 2.3_2_6_6
])
_lowerCAmelCase							:   Union[str, Any]				=				torch.tensor([
    -1.3_5_2_5, -1.9_6_2_8, -0.3_9_5_6, -0.6_8_6_0, 1.4_6_6_4, 1.0_0_1_4, -0.1_2_5_9, -2.7_2_1_2,
    0.7_7_7_2, -1.8_8_1_1, 0.2_9_9_6, 0.4_3_8_8, 0.1_7_0_4, 1.4_0_2_9, 1.1_7_0_1, -0.3_0_2_7,
    -2.2_0_5_3, -2.6_2_8_7, 1.3_3_5_0, 0.8_1_3_1, 0.9_2_7_4, -1.6_2_9_2, -0.6_0_9_8, 2.5_1_3_1,
    1.6_5_0_5, -0.8_9_5_8, -0.9_2_9_8, -0.0_1_5_1, 2.4_2_5_7, 2.3_3_5_5
])
_lowerCAmelCase							:   str				=				torch.tensor([
    -2.0_5_8_5, -2.7_8_9_7, -0.2_8_5_0, -0.8_9_4_0, 1.9_0_5_2, 0.5_7_0_2, 0.6_3_4_5, -3.8_9_5_9,
    1.5_9_3_2, -3.2_3_1_9, 0.1_9_7_4, 0.0_2_8_7, 1.7_5_6_6, 2.6_5_4_3, 0.8_3_8_7, -0.5_3_5_1,
    -3.2_7_3_6, -4.3_3_7_5, 2.9_0_2_9, 1.6_3_9_0, 1.4_6_4_0, -2.1_7_0_1, -1.9_0_1_3, 2.9_3_4_1,
    3.4_9_8_1, -0.6_2_5_5, -1.1_6_4_4, -0.1_5_9_1, 3.7_0_9_7, 3.2_0_6_6
])
_lowerCAmelCase							:   Optional[int]				=				torch.tensor([
    -2.3_1_3_9, -2.5_5_9_4, -0.0_1_9_7, -0.6_7_8_5, 1.7_0_0_1, 1.1_6_0_6, 0.3_0_7_5, -2.1_7_4_0,
    1.8_0_7_1, -2.5_6_3_0, -0.0_9_2_6, -0.3_8_1_1, 1.2_1_1_6, 2.6_2_4_6, 1.2_7_3_1, -0.5_3_9_8,
    -2.8_1_5_3, -3.6_1_4_0, 2.3_8_9_3, 1.3_2_6_2, 1.6_2_5_8, -2.1_8_5_6, -1.3_2_6_7, 2.8_3_9_5,
    2.3_7_7_9, -1.0_6_2_3, -1.2_4_6_8, 0.8_9_5_9, 3.3_3_6_7, 3.2_2_4_3
])
_lowerCAmelCase							:   Union[str, Any]				=				torch.tensor([
    -2.0_6_2_8, -2.7_6_6_7, -0.2_0_8_9, -0.8_2_6_3, 2.0_5_3_9, 0.5_9_9_2, 0.6_4_9_5, -3.8_3_3_6,
    1.6_0_2_5, -3.2_8_1_7, 0.1_7_2_1, -0.0_6_3_3, 1.7_5_1_6, 2.7_0_3_9, 0.8_1_0_0, -0.5_9_0_8,
    -3.2_1_1_3, -4.4_3_4_3, 2.9_2_5_7, 1.3_6_3_2, 1.5_5_6_2, -2.1_4_8_9, -1.9_8_9_4, 3.0_5_6_0,
    3.3_3_9_6, -0.7_3_2_8, -1.0_4_1_7, 0.0_3_8_3, 3.7_0_9_3, 3.2_3_4_3
])
_lowerCAmelCase							:   Any				=				torch.tensor([
    -1.4_5_7_4, -2.0_5_6_9, -0.0_4_7_3, -0.6_1_1_7, 1.4_0_1_8, 0.5_7_6_9, 0.4_1_2_9, -2.7_3_4_4,
    1.2_2_4_1, -2.1_3_9_7, 0.2_0_0_0, 0.3_9_3_7, 0.7_6_1_6, 2.0_4_5_3, 0.7_3_2_4, -0.3_3_9_1,
    -2.1_7_4_6, -2.7_7_4_4, 1.6_9_6_3, 0.6_9_2_1, 1.2_1_8_7, -1.6_1_7_2, -0.8_8_7_7, 2.2_4_3_9,
    1.8_4_7_1, -0.5_8_3_9, -0.5_6_0_5, -0.0_4_6_4, 2.3_2_5_0, 2.1_2_1_9
])
# fmt: on

_lowerCAmelCase							:   Optional[Any]				=				api.list_models(filter="diffusers")
for mod in models:
 if "google" in mod.author or mod.modelId == "CompVis/ldm-celebahq-256":
  _lowerCAmelCase							:   List[Any]				=				"/home/patrick/google_checkpoints/" + mod.modelId.split("/")[-1]

  print(F"Started running {mod.modelId}!!!")

  if mod.modelId.startswith("CompVis"):
   _lowerCAmelCase							:   Tuple				=				UNetaDModel.from_pretrained(local_checkpoint, subfolder="unet")
  else:
   _lowerCAmelCase							:   List[Any]				=				UNetaDModel.from_pretrained(local_checkpoint)

  torch.manual_seed(0)
  random.seed(0)

  _lowerCAmelCase							:   Optional[Any]				=				torch.randn(1, model.config.in_channels, model.config.sample_size, model.config.sample_size)
  _lowerCAmelCase							:   Tuple				=				torch.tensor([10] * noise.shape[0])
  with torch.no_grad():
   _lowerCAmelCase							:   Optional[Any]				=				model(noise, time_step).sample

  assert torch.allclose(
      logits[0, 0, 0, :30], results["_".join("_".join(mod.modelId.split("/")).split("-"))], atol=1E-3
  )
  print(F"{mod.modelId} has passed successfully!!!")
