




'''simple docstring'''




import argparse
from collections import OrderedDict
from pathlib import Path

import torch
from huggingface_hub import hf_hub_download
from PIL import Image
from torchvision.transforms import functional as F

from transformers import DetrImageProcessor, TableTransformerConfig, TableTransformerForObjectDetection
from transformers.utils import logging


logging.set_verbosity_info()
_a   :	Optional[Any] 							=   logging.get_logger(__name__)

# here we list all keys to be renamed (original name on the left, our name on the right)
_a   :	Any 							=   []
for i in range(6):
   # encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
   rename_keys.append(
       (f"""transformer.encoder.layers.{i}.self_attn.out_proj.weight""", f"""encoder.layers.{i}.self_attn.out_proj.weight""")
   )
   rename_keys.append(
       (f"""transformer.encoder.layers.{i}.self_attn.out_proj.bias""", f"""encoder.layers.{i}.self_attn.out_proj.bias""")
   )
   rename_keys.append((f"""transformer.encoder.layers.{i}.linear1.weight""", f"""encoder.layers.{i}.fc1.weight"""))
   rename_keys.append((f"""transformer.encoder.layers.{i}.linear1.bias""", f"""encoder.layers.{i}.fc1.bias"""))
   rename_keys.append((f"""transformer.encoder.layers.{i}.linear2.weight""", f"""encoder.layers.{i}.fc2.weight"""))
   rename_keys.append((f"""transformer.encoder.layers.{i}.linear2.bias""", f"""encoder.layers.{i}.fc2.bias"""))
   rename_keys.append(
       (f"""transformer.encoder.layers.{i}.norm1.weight""", f"""encoder.layers.{i}.self_attn_layer_norm.weight""")
   )
   rename_keys.append((f"""transformer.encoder.layers.{i}.norm1.bias""", f"""encoder.layers.{i}.self_attn_layer_norm.bias"""))
   rename_keys.append((f"""transformer.encoder.layers.{i}.norm2.weight""", f"""encoder.layers.{i}.final_layer_norm.weight"""))
   rename_keys.append((f"""transformer.encoder.layers.{i}.norm2.bias""", f"""encoder.layers.{i}.final_layer_norm.bias"""))
   # decoder layers: 2 times output projection, 2 feedforward neural networks and 3 layernorms
   rename_keys.append(
       (f"""transformer.decoder.layers.{i}.self_attn.out_proj.weight""", f"""decoder.layers.{i}.self_attn.out_proj.weight""")
   )
   rename_keys.append(
       (f"""transformer.decoder.layers.{i}.self_attn.out_proj.bias""", f"""decoder.layers.{i}.self_attn.out_proj.bias""")
   )
   rename_keys.append(
       (
           f"""transformer.decoder.layers.{i}.multihead_attn.out_proj.weight""",
           f"""decoder.layers.{i}.encoder_attn.out_proj.weight""",
       )
   )
   rename_keys.append(
       (
           f"""transformer.decoder.layers.{i}.multihead_attn.out_proj.bias""",
           f"""decoder.layers.{i}.encoder_attn.out_proj.bias""",
       )
   )
   rename_keys.append((f"""transformer.decoder.layers.{i}.linear1.weight""", f"""decoder.layers.{i}.fc1.weight"""))
   rename_keys.append((f"""transformer.decoder.layers.{i}.linear1.bias""", f"""decoder.layers.{i}.fc1.bias"""))
   rename_keys.append((f"""transformer.decoder.layers.{i}.linear2.weight""", f"""decoder.layers.{i}.fc2.weight"""))
   rename_keys.append((f"""transformer.decoder.layers.{i}.linear2.bias""", f"""decoder.layers.{i}.fc2.bias"""))
   rename_keys.append(
       (f"""transformer.decoder.layers.{i}.norm1.weight""", f"""decoder.layers.{i}.self_attn_layer_norm.weight""")
   )
   rename_keys.append((f"""transformer.decoder.layers.{i}.norm1.bias""", f"""decoder.layers.{i}.self_attn_layer_norm.bias"""))
   rename_keys.append(
       (f"""transformer.decoder.layers.{i}.norm2.weight""", f"""decoder.layers.{i}.encoder_attn_layer_norm.weight""")
   )
   rename_keys.append(
       (f"""transformer.decoder.layers.{i}.norm2.bias""", f"""decoder.layers.{i}.encoder_attn_layer_norm.bias""")
   )
   rename_keys.append((f"""transformer.decoder.layers.{i}.norm3.weight""", f"""decoder.layers.{i}.final_layer_norm.weight"""))
   rename_keys.append((f"""transformer.decoder.layers.{i}.norm3.bias""", f"""decoder.layers.{i}.final_layer_norm.bias"""))

# convolutional projection + query embeddings + layernorm of encoder + layernorm of decoder + class and bounding box heads
rename_keys.extend(
    [
        ("input_proj.weight", "input_projection.weight"),
        ("input_proj.bias", "input_projection.bias"),
        ("query_embed.weight", "query_position_embeddings.weight"),
        ("transformer.encoder.norm.weight", "encoder.layernorm.weight"),
        ("transformer.encoder.norm.bias", "encoder.layernorm.bias"),
        ("transformer.decoder.norm.weight", "decoder.layernorm.weight"),
        ("transformer.decoder.norm.bias", "decoder.layernorm.bias"),
        ("class_embed.weight", "class_labels_classifier.weight"),
        ("class_embed.bias", "class_labels_classifier.bias"),
        ("bbox_embed.layers.0.weight", "bbox_predictor.layers.0.weight"),
        ("bbox_embed.layers.0.bias", "bbox_predictor.layers.0.bias"),
        ("bbox_embed.layers.1.weight", "bbox_predictor.layers.1.weight"),
        ("bbox_embed.layers.1.bias", "bbox_predictor.layers.1.bias"),
        ("bbox_embed.layers.2.weight", "bbox_predictor.layers.2.weight"),
        ("bbox_embed.layers.2.bias", "bbox_predictor.layers.2.bias"),
    ]
)

def _lowercase	(      lowerCamelCase__							,    lowerCamelCase__							,    lowerCamelCase__						)					->       List[Any]:





       """simple docstring"""


       __UpperCAmelCase  :				List[str]       				=		state_dict.pop(_lowercase						)
       __UpperCAmelCase  :				str       				=		val

def _lowercase	(      lowerCamelCase__						)					->       str:





       """simple docstring"""


       __UpperCAmelCase  :				List[str]       				=		OrderedDict()
       for key, value in state_dict.items():
              if "backbone.0.body" in key:
                     __UpperCAmelCase  :				List[str]       				=		key.replace("backbone.0.body"							,    "backbone.conv_encoder.model"						)
                     __UpperCAmelCase  :				str       				=		value
              else:
                     __UpperCAmelCase  :				str       				=		value

       return new_state_dict

def _lowercase	(      lowerCamelCase__						)					->       Optional[int]:





       """simple docstring"""


       __UpperCAmelCase  :				Union[str, Any]       				=		''''''

       # first: transformer encoder
       for i in range(6						):
              # read in weights + bias of input projection layer (in PyTorch's MultiHeadAttention, this is a single matrix + bias)
              __UpperCAmelCase  :				Dict       				=		state_dict.pop(f"""{prefix}transformer.encoder.layers.{i}.self_attn.in_proj_weight"""						)
              __UpperCAmelCase  :				List[str]       				=		state_dict.pop(f"""{prefix}transformer.encoder.layers.{i}.self_attn.in_proj_bias"""						)
              # next, add query, keys and values (in that order) to the state dict
              __UpperCAmelCase  :				Tuple       				=		in_proj_weight[:256, :]
              __UpperCAmelCase  :				int       				=		in_proj_bias[:256]
              __UpperCAmelCase  :				Optional[Any]       				=		in_proj_weight[256:512, :]
              __UpperCAmelCase  :				List[str]       				=		in_proj_bias[256:512]
              __UpperCAmelCase  :				Tuple       				=		in_proj_weight[-256:, :]
              __UpperCAmelCase  :				Optional[Any]       				=		in_proj_bias[-256:]
       # next: transformer decoder (which is a bit more complex because it also includes cross-attention)
       for i in range(6						):
              # read in weights + bias of input projection layer of self-attention
              __UpperCAmelCase  :				Optional[int]       				=		state_dict.pop(f"""{prefix}transformer.decoder.layers.{i}.self_attn.in_proj_weight"""						)
              __UpperCAmelCase  :				Dict       				=		state_dict.pop(f"""{prefix}transformer.decoder.layers.{i}.self_attn.in_proj_bias"""						)
              # next, add query, keys and values (in that order) to the state dict
              __UpperCAmelCase  :				str       				=		in_proj_weight[:256, :]
              __UpperCAmelCase  :				List[str]       				=		in_proj_bias[:256]
              __UpperCAmelCase  :				Any       				=		in_proj_weight[256:512, :]
              __UpperCAmelCase  :				Any       				=		in_proj_bias[256:512]
              __UpperCAmelCase  :				List[Any]       				=		in_proj_weight[-256:, :]
              __UpperCAmelCase  :				List[str]       				=		in_proj_bias[-256:]
              # read in weights + bias of input projection layer of cross-attention
              __UpperCAmelCase  :				Any       				=		state_dict.pop(
                  f"""{prefix}transformer.decoder.layers.{i}.multihead_attn.in_proj_weight"""						)
              __UpperCAmelCase  :				List[str]       				=		state_dict.pop(f"""{prefix}transformer.decoder.layers.{i}.multihead_attn.in_proj_bias"""						)
              # next, add query, keys and values (in that order) of cross-attention to the state dict
              __UpperCAmelCase  :				Tuple       				=		in_proj_weight_cross_attn[:256, :]
              __UpperCAmelCase  :				Optional[Any]       				=		in_proj_bias_cross_attn[:256]
              __UpperCAmelCase  :				Dict       				=		in_proj_weight_cross_attn[256:512, :]
              __UpperCAmelCase  :				Tuple       				=		in_proj_bias_cross_attn[256:512]
              __UpperCAmelCase  :				Optional[int]       				=		in_proj_weight_cross_attn[-256:, :]
              __UpperCAmelCase  :				Dict       				=		in_proj_bias_cross_attn[-256:]

def _lowercase	(      lowerCamelCase__							,    lowerCamelCase__						)					->       int:





       """simple docstring"""


       __UpperCAmelCase  :				List[str]       				=		image.size
       __UpperCAmelCase  :				Tuple       				=		max(_lowercase							,    _lowercase						)
       __UpperCAmelCase  :				Tuple       				=		800 if '''detection''' in checkpoint_url else 1000
       __UpperCAmelCase  :				str       				=		target_max_size / current_max_size
       __UpperCAmelCase  :				List[str]       				=		image.resize((int(round(scale * width						)						), int(round(scale * height						)						))						)

       return resized_image

def _lowercase	(      lowerCamelCase__						)					->       Dict:





       """simple docstring"""


       __UpperCAmelCase  :				Optional[Any]       				=		F.to_tensor(_lowercase						)
       __UpperCAmelCase  :				str       				=		F.normalize(_lowercase							,    mean=[0.485, 0.456, 0.406]							,    std=[0.229, 0.224, 0.225]						)
       return image

@torch.no_grad()
def _lowercase	(      lowerCamelCase__							,    lowerCamelCase__							,    lowerCamelCase__						)					->       List[str]:





       """simple docstring"""


       logger.info("Converting model..."						)

       # load original state dict
       __UpperCAmelCase  :				List[Any]       				=		torch.hub.load_state_dict_from_url(_lowercase							,    map_location="cpu"						)
       # rename keys
       for src, dest in rename_keys:
              rename_key(_lowercase							,    _lowercase							,    _lowercase						)
       __UpperCAmelCase  :				List[Any]       				=		rename_backbone_keys(_lowercase						)
       # query, key and value matrices need special treatment
       read_in_q_k_v(_lowercase						)
       # important: we need to prepend a prefix to each of the base model keys as the head models use different attributes for them
       __UpperCAmelCase  :				Any       				=		'''model.'''
       for key in state_dict.copy().keys():
              if not key.startswith("class_labels_classifier"						) and not key.startswith("bbox_predictor"						):
                     __UpperCAmelCase  :				int       				=		state_dict.pop(_lowercase						)
                     __UpperCAmelCase  :				int       				=		val
    # create HuggingFace model and load state dict
       __UpperCAmelCase  :				Optional[int]       				=		TableTransformerConfig(
           backbone="resnet18"							,    mask_loss_coefficient=1							,    dice_loss_coefficient=1							,    ce_loss_coefficient=1							,    bbox_loss_coefficient=5							,    giou_loss_coefficient=2							,    eos_coefficient=0.4							,    class_cost=1							,    bbox_cost=5							,    giou_cost=2							,    )

       if "detection" in checkpoint_url:
              __UpperCAmelCase  :				Optional[Any]       				=		15
              __UpperCAmelCase  :				int       				=		2
              __UpperCAmelCase  :				Optional[int]       				=		{0: '''table''', 1: '''table rotated'''}
              __UpperCAmelCase  :				Dict       				=		idalabel
              __UpperCAmelCase  :				int       				=		{v: k for k, v in idalabel.items()}
       else:
              __UpperCAmelCase  :				List[str]       				=		125
              __UpperCAmelCase  :				Optional[int]       				=		6
              __UpperCAmelCase  :				Union[str, Any]       				=		{
                  0: '''table''',
                  1: '''table column''',
                  2: '''table row''',
                  3: '''table column header''',
                  4: '''table projected row header''',
                  5: '''table spanning cell''',
              }
              __UpperCAmelCase  :				Optional[int]       				=		idalabel
              __UpperCAmelCase  :				List[Any]       				=		{v: k for k, v in idalabel.items()}

       __UpperCAmelCase  :				Union[str, Any]       				=		DetrImageProcessor(
           format="coco_detection"							,    max_size=800 if "detection" in checkpoint_url else 1000						)
       __UpperCAmelCase  :				int       				=		TableTransformerForObjectDetection(_lowercase						)
       model.load_state_dict(_lowercase						)
       model.eval()

       # verify our conversion
       __UpperCAmelCase  :				List[Any]       				=		'''example_pdf.png''' if '''detection''' in checkpoint_url else '''example_table.png'''
       __UpperCAmelCase  :				int       				=		hf_hub_download(repo_id="nielsr/example-pdf"							,    repo_type="dataset"							,    filename=_lowercase						)
       __UpperCAmelCase  :				List[Any]       				=		Image.open(_lowercase						).convert("RGB"						)
       __UpperCAmelCase  :				List[str]       				=		normalize(resize(_lowercase							,    _lowercase						)						).unsqueeze(0						)

       __UpperCAmelCase  :				List[Any]       				=		model(_lowercase						)

       if "detection" in checkpoint_url:
              __UpperCAmelCase  :				Dict       				=		(1, 15, 3)
              __UpperCAmelCase  :				Optional[Any]       				=		torch.tensor(
                  [[-6.7897, -16.9985, 6.7937], [-8.0186, -22.2192, 6.9677], [-7.3117, -21.0708, 7.4055]]						)
              __UpperCAmelCase  :				Optional[int]       				=		torch.tensor([[0.4867, 0.1767, 0.6732], [0.6718, 0.4479, 0.3830], [0.4716, 0.1760, 0.6364]]						)

       else:
              __UpperCAmelCase  :				int       				=		(1, 125, 7)
              __UpperCAmelCase  :				Tuple       				=		torch.tensor(
                  [[-18.1430, -8.3214, 4.8274], [-18.4685, -7.1361, -4.2667], [-26.3693, -9.3429, -4.9962]]						)
              __UpperCAmelCase  :				Optional[Any]       				=		torch.tensor([[0.4983, 0.5595, 0.9440], [0.4916, 0.6315, 0.5954], [0.6108, 0.8637, 0.1135]]						)

       assert outputs.logits.shape == expected_shape
       assert torch.allclose(outputs.logits[0, :3, :3]							,    _lowercase							,    atol=1e-4						)
       assert torch.allclose(outputs.pred_boxes[0, :3, :3]							,    _lowercase							,    atol=1e-4						)
       print("Looks ok!"						)

       if pytorch_dump_folder_path is not None:
              # Save model and image processor
              logger.info(f"""Saving PyTorch model and image processor to {pytorch_dump_folder_path}..."""						)
              Path(_lowercase						).mkdir(exist_ok=_lowercase						)
              model.save_pretrained(_lowercase						)
              image_processor.save_pretrained(_lowercase						)

       if push_to_hub:
              # Push model to HF hub
              logger.info("Pushing model to the hub..."						)
              __UpperCAmelCase  :				Dict       				=		(
                  '''microsoft/table-transformer-detection'''
                  if '''detection''' in checkpoint_url
                  else '''microsoft/table-transformer-structure-recognition'''
              )
              model.push_to_hub(_lowercase						)
              image_processor.push_to_hub(_lowercase						)


if __name__ == "__main__":
   _a   :	str 							=   argparse.ArgumentParser()

   parser.add_argument(
       "--checkpoint_url",
       default="https://pubtables1m.blob.core.windows.net/model/pubtables1m_detection_detr_r18.pth",
       type=str,
       choices=[
           "https://pubtables1m.blob.core.windows.net/model/pubtables1m_detection_detr_r18.pth",
           "https://pubtables1m.blob.core.windows.net/model/pubtables1m_structure_detr_r18.pth",
       ],
       help="URL of the Table Transformer checkpoint you\'d like to convert.",
   )
   parser.add_argument(
       "--pytorch_dump_folder_path", default=None, type=str, help="Path to the folder to output PyTorch model."
   )
   parser.add_argument(
       "--push_to_hub", action="store_true", help="Whether or not to push the converted model to the 🤗 hub."
   )
   _a   :	Optional[int] 							=   parser.parse_args()
   convert_table_transformer_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path, args.push_to_hub)







'''simple docstring'''




import random

import torch
from huggingface_hub import HfApi

from diffusers import UNetaDModel


_a   :	Union[str, Any] 							=   HfApi()

_a   :	int 							=   {}
# fmt: off
_a   :	Optional[int] 							=   torch.tensor([
    -0.7_515, -1.6_883, 0.2_420, 0.0_300, 0.6_347, 1.3_433, -1.1_743, -3.7_467,
    1.2_342, -2.2_485, 0.4_636, 0.8_076, -0.7_991, 0.3_969, 0.8_498, 0.9_189,
    -1.8_887, -3.3_522, 0.7_639, 0.2_040, 0.6_271, -2.7_148, -1.6_316, 3.0_839,
    0.3_186, 0.2_721, -0.9_759, -1.2_461, 2.6_257, 1.3_557
])
_a   :	Optional[Any] 							=   torch.tensor([
    -2.3_639, -2.5_344, 0.0_054, -0.6_674, 1.5_990, 1.0_158, 0.3_124, -2.1_436,
    1.8_795, -2.5_429, -0.1_566, -0.3_973, 1.2_490, 2.6_447, 1.2_283, -0.5_208,
    -2.8_154, -3.5_119, 2.3_838, 1.2_033, 1.7_201, -2.1_256, -1.4_576, 2.7_948,
    2.4_204, -0.9_752, -1.2_546, 0.8_027, 3.2_758, 3.1_365
])
_a   :	int 							=   torch.tensor([
    -0.6_531, -0.6_891, -0.3_172, -0.5_375, -0.9_140, -0.5_367, -0.1_175, -0.7_869,
    -0.3_808, -0.4_513, -0.2_098, -0.0_083, 0.3_183, 0.5_140, 0.2_247, -0.1_304,
    -0.1_302, -0.2_802, -0.2_084, -0.2_025, -0.4_967, -0.4_873, -0.0_861, 0.6_925,
    0.0_250, 0.1_290, -0.1_543, 0.6_316, 1.0_460, 1.4_943
])
_a   :	str 							=   torch.tensor([
    0.0_911, 0.1_107, 0.0_182, 0.0_435, -0.0_805, -0.0_608, 0.0_381, 0.2_172,
    -0.0_280, 0.1_327, -0.0_299, -0.0_255, -0.0_050, -0.1_170, -0.1_046, 0.0_309,
    0.1_367, 0.1_728, -0.0_533, -0.0_748, -0.0_534, 0.1_624, 0.0_384, -0.1_805,
    -0.0_707, 0.0_642, 0.0_220, -0.0_134, -0.1_333, -0.1_505
])
_a   :	Union[str, Any] 							=   torch.tensor([
    0.1_321, 0.1_337, 0.0_440, 0.0_622, -0.0_591, -0.0_370, 0.0_503, 0.2_133,
    -0.0_177, 0.1_415, -0.0_116, -0.0_112, 0.0_044, -0.0_980, -0.0_789, 0.0_395,
    0.1_502, 0.1_785, -0.0_488, -0.0_514, -0.0_404, 0.1_539, 0.0_454, -0.1_559,
    -0.0_665, 0.0_659, 0.0_383, -0.0_005, -0.1_266, -0.1_386
])
_a   :	Any 							=   torch.tensor([
    0.1_154, 0.1_218, 0.0_307, 0.0_526, -0.0_711, -0.0_541, 0.0_366, 0.2_078,
    -0.0_267, 0.1_317, -0.0_226, -0.0_193, -0.0_014, -0.1_055, -0.0_902, 0.0_330,
    0.1_391, 0.1_709, -0.0_562, -0.0_693, -0.0_560, 0.1_482, 0.0_381, -0.1_683,
    -0.0_681, 0.0_661, 0.0_331, -0.0_046, -0.1_268, -0.1_431
])
_a   :	List[Any] 							=   torch.tensor([
    0.1_192, 0.1_240, 0.0_414, 0.0_606, -0.0_557, -0.0_412, 0.0_430, 0.2_042,
    -0.0_200, 0.1_385, -0.0_115, -0.0_132, 0.0_017, -0.0_965, -0.0_802, 0.0_398,
    0.1_433, 0.1_747, -0.0_458, -0.0_533, -0.0_407, 0.1_545, 0.0_419, -0.1_574,
    -0.0_645, 0.0_626, 0.0_341, -0.0_010, -0.1_199, -0.1_390
])
_a   :	Optional[int] 							=   torch.tensor([
    0.1_075, 0.1_074, 0.0_205, 0.0_431, -0.0_774, -0.0_607, 0.0_298, 0.2_042,
    -0.0_320, 0.1_267, -0.0_281, -0.0_250, -0.0_064, -0.1_091, -0.0_946, 0.0_290,
    0.1_328, 0.1_650, -0.0_580, -0.0_738, -0.0_586, 0.1_440, 0.0_337, -0.1_746,
    -0.0_712, 0.0_605, 0.0_250, -0.0_099, -0.1_316, -0.1_473
])
_a   :	Tuple 							=   torch.tensor([
    -1.4_572, -2.0_481, -0.0_414, -0.6_005, 1.4_136, 0.5_848, 0.4_028, -2.7_330,
    1.2_212, -2.1_228, 0.2_155, 0.4_039, 0.7_662, 2.0_535, 0.7_477, -0.3_243,
    -2.1_758, -2.7_648, 1.6_947, 0.7_026, 1.2_338, -1.6_078, -0.8_682, 2.2_810,
    1.8_574, -0.5_718, -0.5_586, -0.0_186, 2.3_415, 2.1_251])
_a   :	List[Any] 							=   torch.tensor([
    -1.3_690, -1.9_720, -0.4_090, -0.6_966, 1.4_660, 0.9_938, -0.1_385, -2.7_324,
    0.7_736, -1.8_917, 0.2_923, 0.4_293, 0.1_693, 1.4_112, 1.1_887, -0.3_181,
    -2.2_160, -2.6_381, 1.3_170, 0.8_163, 0.9_240, -1.6_544, -0.6_099, 2.5_259,
    1.6_430, -0.9_090, -0.9_392, -0.0_126, 2.4_268, 2.3_266
])
_a   :	Optional[Any] 							=   torch.tensor([
    -1.3_525, -1.9_628, -0.3_956, -0.6_860, 1.4_664, 1.0_014, -0.1_259, -2.7_212,
    0.7_772, -1.8_811, 0.2_996, 0.4_388, 0.1_704, 1.4_029, 1.1_701, -0.3_027,
    -2.2_053, -2.6_287, 1.3_350, 0.8_131, 0.9_274, -1.6_292, -0.6_098, 2.5_131,
    1.6_505, -0.8_958, -0.9_298, -0.0_151, 2.4_257, 2.3_355
])
_a   :	Union[str, Any] 							=   torch.tensor([
    -2.0_585, -2.7_897, -0.2_850, -0.8_940, 1.9_052, 0.5_702, 0.6_345, -3.8_959,
    1.5_932, -3.2_319, 0.1_974, 0.0_287, 1.7_566, 2.6_543, 0.8_387, -0.5_351,
    -3.2_736, -4.3_375, 2.9_029, 1.6_390, 1.4_640, -2.1_701, -1.9_013, 2.9_341,
    3.4_981, -0.6_255, -1.1_644, -0.1_591, 3.7_097, 3.2_066
])
_a   :	Optional[int] 							=   torch.tensor([
    -2.3_139, -2.5_594, -0.0_197, -0.6_785, 1.7_001, 1.1_606, 0.3_075, -2.1_740,
    1.8_071, -2.5_630, -0.0_926, -0.3_811, 1.2_116, 2.6_246, 1.2_731, -0.5_398,
    -2.8_153, -3.6_140, 2.3_893, 1.3_262, 1.6_258, -2.1_856, -1.3_267, 2.8_395,
    2.3_779, -1.0_623, -1.2_468, 0.8_959, 3.3_367, 3.2_243
])
_a   :	Union[str, Any] 							=   torch.tensor([
    -2.0_628, -2.7_667, -0.2_089, -0.8_263, 2.0_539, 0.5_992, 0.6_495, -3.8_336,
    1.6_025, -3.2_817, 0.1_721, -0.0_633, 1.7_516, 2.7_039, 0.8_100, -0.5_908,
    -3.2_113, -4.4_343, 2.9_257, 1.3_632, 1.5_562, -2.1_489, -1.9_894, 3.0_560,
    3.3_396, -0.7_328, -1.0_417, 0.0_383, 3.7_093, 3.2_343
])
_a   :	str 							=   torch.tensor([
    -1.4_574, -2.0_569, -0.0_473, -0.6_117, 1.4_018, 0.5_769, 0.4_129, -2.7_344,
    1.2_241, -2.1_397, 0.2_000, 0.3_937, 0.7_616, 2.0_453, 0.7_324, -0.3_391,
    -2.1_746, -2.7_744, 1.6_963, 0.6_921, 1.2_187, -1.6_172, -0.8_877, 2.2_439,
    1.8_471, -0.5_839, -0.5_605, -0.0_464, 2.3_250, 2.1_219
])
# fmt: on

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

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

      if mod.modelId.startswith("CompVis"):
         _a   :	int 							=   UNetaDModel.from_pretrained(local_checkpoint, subfolder="unet")
      else:
         _a   :	Optional[int] 							=   UNetaDModel.from_pretrained(local_checkpoint)

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

      _a   :	str 							=   torch.randn(1, model.config.in_channels, model.config.sample_size, model.config.sample_size)
      _a   :	str 							=   torch.tensor([10] * noise.shape[0])
      with torch.no_grad():
         _a   :	str 							=   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!!!""")






