




import gc
import unittest

import torch
from parameterized import parameterized

from diffusers import AutoencoderKL
from diffusers.utils import floats_tensor, load_hf_numpy, require_torch_gpu, slow, torch_all_close, torch_device
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin, UNetTesterMixin


enable_full_determinism()






class    UpperCAmelCase__	( A_       ,  A_       ,  unittest.TestCase		):



    """simple docstring"""






    UpperCAmelCase__						:		Optional[Any]              =  AutoencoderKL
    UpperCAmelCase__						:		Tuple              =  "sample"
    UpperCAmelCase__						:		Optional[int]              =  1e-2






    @property
    def _a (      self   )  ->	int:
         __UpperCamelCase   							=4
         __UpperCamelCase   							=3
         __UpperCamelCase   							=(32, 32)

         __UpperCamelCase   							=floats_tensor((batch_size, num_channels) + sizes   ).to(A_   )

         return {"sample": image}






    @property
    def _a (      self   )  ->	Any:
         return (3, 32, 32)






    @property
    def _a (      self   )  ->	Union[str, Any]:
         return (3, 32, 32)






    def _a (      self   )  ->	Union[str, Any]:
         __UpperCamelCase   							={
             'block_out_channels': [32, 64],
             'in_channels': 3,
             'out_channels': 3,
             'down_block_types': ['DownEncoderBlock2D', 'DownEncoderBlock2D'],
             'up_block_types': ['UpDecoderBlock2D', 'UpDecoderBlock2D'],
             'latent_channels': 4,
         }
         __UpperCamelCase   							=self.dummy_input
         return init_dict, inputs_dict






    def _a (      self   )  ->	List[Any]:
         pass






    def _a (      self   )  ->	Union[str, Any]:
         pass






    @unittest.skipIf(torch_device == 'mps' ,       'Gradient checkpointing skipped on MPS'   )
    def _a (      self   )  ->	str:
         # enable deterministic behavior for gradient checkpointing
         __UpperCamelCase							,				__UpperCamelCase   							=self.prepare_init_args_and_inputs_for_common()
         __UpperCamelCase   							=self.model_class(**A_   )
         model.to(A_   )

         assert not model.is_gradient_checkpointing and model.training

         __UpperCamelCase   							=model(**A_   ).sample
         # run the backwards pass on the model. For backwards pass, for simplicity purpose,
         # we won't calculate the loss and rather backprop on out.sum()
         model.zero_grad()

         __UpperCamelCase   							=torch.randn_like(A_   )
         __UpperCamelCase   							=(out - labels).mean()
         loss.backward()

         # re-instantiate the model now enabling gradient checkpointing
         __UpperCamelCase   							=self.model_class(**A_   )
         # clone model
         model_a.load_state_dict(model.state_dict()   )
         model_a.to(A_   )
         model_a.enable_gradient_checkpointing()

         assert model_a.is_gradient_checkpointing and model_a.training

         __UpperCamelCase   							=model_a(**A_   ).sample
         # run the backwards pass on the model. For backwards pass, for simplicity purpose,
         # we won't calculate the loss and rather backprop on out.sum()
         model_a.zero_grad()
         __UpperCamelCase   							=(out_a - labels).mean()
         loss_a.backward()

         # compare the output and parameters gradients
         self.assertTrue((loss - loss_a).abs() < 1E-5   )
         __UpperCamelCase   							=dict(model.named_parameters()   )
         __UpperCamelCase   							=dict(model_a.named_parameters()   )
         for name, param in named_params.items():
              self.assertTrue(torch_all_close(param.grad.data ,       named_params_a[name].grad.data ,       atol=5E-5   )   )






    def _a (      self   )  ->	Union[str, Any]:
         __UpperCamelCase							,				__UpperCamelCase   							=AutoencoderKL.from_pretrained('fusing/autoencoder-kl-dummy' ,       output_loading_info=A_   )
         self.assertIsNotNone(A_   )
         self.assertEqual(len(loading_info['missing_keys']   ) ,       0   )

         model.to(A_   )
         __UpperCamelCase   							=model(**self.dummy_input   )

         assert image is not None, "Make sure output is not None"






    def _a (      self   )  ->	str:
         __UpperCamelCase   							=AutoencoderKL.from_pretrained('fusing/autoencoder-kl-dummy'   )
         __UpperCamelCase   							=model.to(A_   )
         model.eval()

         if torch_device == "mps":
              __UpperCamelCase   							=torch.manual_seed(0   )
         else:
              __UpperCamelCase   							=torch.Generator(device=A_   ).manual_seed(0   )

         __UpperCamelCase   							=torch.randn(
             1 ,       model.config.in_channels ,       model.config.sample_size ,       model.config.sample_size ,       generator=torch.manual_seed(0   ) ,       )
         __UpperCamelCase   							=image.to(A_   )
         with torch.no_grad():
              __UpperCamelCase   							=model(A_ ,       sample_posterior=A_ ,       generator=A_   ).sample

         __UpperCamelCase   							=output[0, -1, -3:, -3:].flatten().cpu()

         # Since the VAE Gaussian prior's generator is seeded on the appropriate device,
         # the expected output slices are not the same for CPU and GPU.
         if torch_device == "mps":
              __UpperCamelCase   							=torch.tensor(
                  [
                      -4.0078E-01,
                      -3.8323E-04,
                      -1.2681E-01,
                      -1.1462E-01,
                      2.0095E-01,
                      1.0893E-01,
                      -8.8247E-02,
                      -3.0361E-01,
                      -9.8644E-03,
                  ]   )
         elif torch_device == "cpu":
              __UpperCamelCase   							=torch.tensor(
                  [-0.1352, 0.0878, 0.0419, -0.0818, -0.1069, 0.0688, -0.1458, -0.4446, -0.0026]   )
         else:
              __UpperCamelCase   							=torch.tensor(
                  [-0.2421, 0.4642, 0.2507, -0.0438, 0.0682, 0.3160, -0.2018, -0.0727, 0.2485]   )

         self.assertTrue(torch_all_close(A_ ,       A_ ,       rtol=1E-2   )   )






@slow
class    UpperCAmelCase__	( unittest.TestCase		):



    """simple docstring"""






    def _a (      self ,       A_ ,       A_   )  ->	Optional[int]:
         return f'gaussian_noise_s={seed}_shape={"_".join([str(A_   ) for s in shape]   )}.npy'






    def _a (      self   )  ->	Union[str, Any]:
         # clean up the VRAM after each test
         super().tearDown()
         gc.collect()
         torch.cuda.empty_cache()






    def _a (      self ,       A_=0 ,       A_=(4, 3, 512, 512) ,       A_=False   )  ->	List[str]:
         __UpperCamelCase   							=torch.floataa if fpaa else torch.floataa
         __UpperCamelCase   							=torch.from_numpy(load_hf_numpy(self.get_file_format(A_ ,       A_   )   )   ).to(A_   ).to(A_   )
         return image






    def _a (      self ,       A_="CompVis/stable-diffusion-v1-4" ,       A_=False   )  ->	Dict:
         __UpperCamelCase   							='fp16' if fpaa else None
         __UpperCamelCase   							=torch.floataa if fpaa else torch.floataa

         __UpperCamelCase   							=AutoencoderKL.from_pretrained(
             A_ ,       subfolder='vae' ,       torch_dtype=A_ ,       revision=A_ ,       )
         model.to(A_   ).eval()

         return model






    def _a (      self ,       A_=0   )  ->	Any:
         if torch_device == "mps":
              return torch.manual_seed(A_   )
         return torch.Generator(device=A_   ).manual_seed(A_   )






    @parameterized.expand(
        [
            # fmt: off
            [33, [-0.1603, 0.9878, -0.0495, -0.0790, -0.2709, 0.8375, -0.2060, -0.0824], [-0.2395, 0.0098, 0.0102, -0.0709, -0.2840, -0.0274, -0.0718, -0.1824]],
            [47, [-0.2376, 0.1168, 0.1332, -0.4840, -0.2508, -0.0791, -0.0493, -0.4089], [0.0350, 0.0847, 0.0467, 0.0344, -0.0842, -0.0547, -0.0633, -0.1131]],
            # fmt: on
        ]   )
    def _a (      self ,       A_ ,       A_ ,       A_   )  ->	List[str]:
         __UpperCamelCase   							=self.get_sd_vae_model()
         __UpperCamelCase   							=self.get_sd_image(A_   )
         __UpperCamelCase   							=self.get_generator(A_   )

         with torch.no_grad():
              __UpperCamelCase   							=model(A_ ,       generator=A_ ,       sample_posterior=A_   ).sample

         assert sample.shape == image.shape

         __UpperCamelCase   							=sample[-1, -2:, -2:, :2].flatten().float().cpu()
         __UpperCamelCase   							=torch.tensor(expected_slice_mps if torch_device == 'mps' else expected_slice   )

         assert torch_all_close(A_ ,       A_ ,       atol=3E-3   )






    @parameterized.expand(
        [
            # fmt: off
            [33, [-0.0513, 0.0289, 1.3799, 0.2166, -0.2573, -0.0871, 0.5103, -0.0999]],
            [47, [-0.4128, -0.1320, -0.3704, 0.1965, -0.4116, -0.2332, -0.3340, 0.2247]],
            # fmt: on
        ]   )
    @require_torch_gpu
    def _a (      self ,       A_ ,       A_   )  ->	int:
         __UpperCamelCase   							=self.get_sd_vae_model(fpaa=A_   )
         __UpperCamelCase   							=self.get_sd_image(A_ ,       fpaa=A_   )
         __UpperCamelCase   							=self.get_generator(A_   )

         with torch.no_grad():
              __UpperCamelCase   							=model(A_ ,       generator=A_ ,       sample_posterior=A_   ).sample

         assert sample.shape == image.shape

         __UpperCamelCase   							=sample[-1, -2:, :2, -2:].flatten().float().cpu()
         __UpperCamelCase   							=torch.tensor(A_   )

         assert torch_all_close(A_ ,       A_ ,       atol=1E-2   )






    @parameterized.expand(
        [
            # fmt: off
            [33, [-0.1609, 0.9866, -0.0487, -0.0777, -0.2716, 0.8368, -0.2055, -0.0814], [-0.2395, 0.0098, 0.0102, -0.0709, -0.2840, -0.0274, -0.0718, -0.1824]],
            [47, [-0.2377, 0.1147, 0.1333, -0.4841, -0.2506, -0.0805, -0.0491, -0.4085], [0.0350, 0.0847, 0.0467, 0.0344, -0.0842, -0.0547, -0.0633, -0.1131]],
            # fmt: on
        ]   )
    def _a (      self ,       A_ ,       A_ ,       A_   )  ->	Optional[Any]:
         __UpperCamelCase   							=self.get_sd_vae_model()
         __UpperCamelCase   							=self.get_sd_image(A_   )

         with torch.no_grad():
              __UpperCamelCase   							=model(A_   ).sample

         assert sample.shape == image.shape

         __UpperCamelCase   							=sample[-1, -2:, -2:, :2].flatten().float().cpu()
         __UpperCamelCase   							=torch.tensor(expected_slice_mps if torch_device == 'mps' else expected_slice   )

         assert torch_all_close(A_ ,       A_ ,       atol=3E-3   )






    @parameterized.expand(
        [
            # fmt: off
            [13, [-0.2051, -0.1803, -0.2311, -0.2114, -0.3292, -0.3574, -0.2953, -0.3323]],
            [37, [-0.2632, -0.2625, -0.2199, -0.2741, -0.4539, -0.4990, -0.3720, -0.4925]],
            # fmt: on
        ]   )
    @require_torch_gpu
    def _a (      self ,       A_ ,       A_   )  ->	str:
         __UpperCamelCase   							=self.get_sd_vae_model()
         __UpperCamelCase   							=self.get_sd_image(A_ ,       shape=(3, 4, 64, 64)   )

         with torch.no_grad():
              __UpperCamelCase   							=model.decode(A_   ).sample

         assert list(sample.shape   ) == [3, 3, 512, 512]

         __UpperCamelCase   							=sample[-1, -2:, :2, -2:].flatten().cpu()
         __UpperCamelCase   							=torch.tensor(A_   )

         assert torch_all_close(A_ ,       A_ ,       atol=1E-3   )






    @parameterized.expand(
        [
            # fmt: off
            [27, [-0.0369, 0.0207, -0.0776, -0.0682, -0.1747, -0.1930, -0.1465, -0.2039]],
            [16, [-0.1628, -0.2134, -0.2747, -0.2642, -0.3774, -0.4404, -0.3687, -0.4277]],
            # fmt: on
        ]   )
    @require_torch_gpu
    def _a (      self ,       A_ ,       A_   )  ->	List[str]:
         __UpperCamelCase   							=self.get_sd_vae_model(fpaa=A_   )
         __UpperCamelCase   							=self.get_sd_image(A_ ,       shape=(3, 4, 64, 64) ,       fpaa=A_   )

         with torch.no_grad():
              __UpperCamelCase   							=model.decode(A_   ).sample

         assert list(sample.shape   ) == [3, 3, 512, 512]

         __UpperCamelCase   							=sample[-1, -2:, :2, -2:].flatten().float().cpu()
         __UpperCamelCase   							=torch.tensor(A_   )

         assert torch_all_close(A_ ,       A_ ,       atol=5E-3   )






    @parameterized.expand([(13,), (16,), (27,)]   )
    @require_torch_gpu
    @unittest.skipIf(not is_xformers_available() ,       reason='xformers is not required when using PyTorch 2.0.'   )
    def _a (      self ,       A_   )  ->	List[Any]:
         __UpperCamelCase   							=self.get_sd_vae_model(fpaa=A_   )
         __UpperCamelCase   							=self.get_sd_image(A_ ,       shape=(3, 4, 64, 64) ,       fpaa=A_   )

         with torch.no_grad():
              __UpperCamelCase   							=model.decode(A_   ).sample

         model.enable_xformers_memory_efficient_attention()
         with torch.no_grad():
              __UpperCamelCase   							=model.decode(A_   ).sample

         assert list(sample.shape   ) == [3, 3, 512, 512]

         assert torch_all_close(A_ ,       A_ ,       atol=1E-1   )






    @parameterized.expand([(13,), (16,), (37,)]   )
    @require_torch_gpu
    @unittest.skipIf(not is_xformers_available() ,       reason='xformers is not required when using PyTorch 2.0.'   )
    def _a (      self ,       A_   )  ->	List[str]:
         __UpperCamelCase   							=self.get_sd_vae_model()
         __UpperCamelCase   							=self.get_sd_image(A_ ,       shape=(3, 4, 64, 64)   )

         with torch.no_grad():
              __UpperCamelCase   							=model.decode(A_   ).sample

         model.enable_xformers_memory_efficient_attention()
         with torch.no_grad():
              __UpperCamelCase   							=model.decode(A_   ).sample

         assert list(sample.shape   ) == [3, 3, 512, 512]

         assert torch_all_close(A_ ,       A_ ,       atol=1E-2   )






    @parameterized.expand(
        [
            # fmt: off
            [33, [-0.3001, 0.0918, -2.6984, -3.9720, -3.2099, -5.0353, 1.7338, -0.2065, 3.4267]],
            [47, [-1.5030, -4.3871, -6.0355, -9.1157, -1.6661, -2.7853, 2.1607, -5.0823, 2.5633]],
            # fmt: on
        ]   )
    def _a (      self ,       A_ ,       A_   )  ->	Optional[int]:
         __UpperCamelCase   							=self.get_sd_vae_model()
         __UpperCamelCase   							=self.get_sd_image(A_   )
         __UpperCamelCase   							=self.get_generator(A_   )

         with torch.no_grad():
              __UpperCamelCase   							=model.encode(A_   ).latent_dist
              __UpperCamelCase   							=dist.sample(generator=A_   )

         assert list(sample.shape   ) == [image.shape[0], 4] + [i // 8 for i in image.shape[2:]]

         __UpperCamelCase   							=sample[0, -1, -3:, -3:].flatten().cpu()
         __UpperCamelCase   							=torch.tensor(A_   )

         __UpperCamelCase   							=3E-3 if torch_device != 'mps' else 1E-2
         assert torch_all_close(A_ ,       A_ ,       atol=A_   )









import argparse
import json
from pathlib import Path

import requests
import torch
from huggingface_hub import cached_download, hf_hub_download, hf_hub_url
from PIL import Image

from transformers import DetaConfig, DetaForObjectDetection, DetaImageProcessor, SwinConfig
from transformers.utils import logging


logging.set_verbosity_info()
_A			=						logging.get_logger(__name__)

def 						_UpperCAmelCase    (      SCREAMING_SNAKE_CASE__      : Optional[Any]    ):
     __UpperCamelCase   							=SwinConfig(
         embed_dim=1_92      ,					depths=(2, 2, 18, 2)      ,					num_heads=(6, 12, 24, 48)      ,					window_size=12      ,					out_features=['stage2', 'stage3', 'stage4']      ,					)

     __UpperCamelCase   							=DetaConfig(
         backbone_config=SCREAMING_SNAKE_CASE__      ,					num_queries=9_00      ,					encoder_ffn_dim=20_48      ,					decoder_ffn_dim=20_48      ,					num_feature_levels=5      ,					assign_first_stage=SCREAMING_SNAKE_CASE__      ,					with_box_refine=SCREAMING_SNAKE_CASE__      ,					two_stage=SCREAMING_SNAKE_CASE__      ,					)

     # set labels
     __UpperCamelCase   							='huggingface/label-files'
     if "o365" in model_name:
          __UpperCamelCase   							=3_66
          __UpperCamelCase   							='object365-id2label.json'
     else:
          __UpperCamelCase   							=91
          __UpperCamelCase   							='coco-detection-id2label.json'

     __UpperCamelCase   							=num_labels
     __UpperCamelCase   							=json.load(open(cached_download(hf_hub_url(SCREAMING_SNAKE_CASE__      ,					SCREAMING_SNAKE_CASE__      ,					repo_type='dataset'    )    )      ,					'r'    )    )
     __UpperCamelCase   							={int(SCREAMING_SNAKE_CASE__    ): v for k, v in idalabel.items()}
     __UpperCamelCase   							=idalabel
     __UpperCamelCase   							={v: k for k, v in idalabel.items()}

     return config

def 						_UpperCAmelCase    (      SCREAMING_SNAKE_CASE__      : Optional[Any]    ):
     __UpperCamelCase   							=[]

     # stem
     # fmt: off
     rename_keys.append(('backbone.0.body.patch_embed.proj.weight', 'model.backbone.model.embeddings.patch_embeddings.projection.weight')    )
     rename_keys.append(('backbone.0.body.patch_embed.proj.bias', 'model.backbone.model.embeddings.patch_embeddings.projection.bias')    )
     rename_keys.append(('backbone.0.body.patch_embed.norm.weight', 'model.backbone.model.embeddings.norm.weight')    )
     rename_keys.append(('backbone.0.body.patch_embed.norm.bias', 'model.backbone.model.embeddings.norm.bias')    )
     # stages
     for i in range(len(config.backbone_config.depths    )    ):
          for j in range(config.backbone_config.depths[i]    ):
               rename_keys.append((F'backbone.0.body.layers.{i}.blocks.{j}.norm1.weight', F'model.backbone.model.encoder.layers.{i}.blocks.{j}.layernorm_before.weight')    )
               rename_keys.append((F'backbone.0.body.layers.{i}.blocks.{j}.norm1.bias', F'model.backbone.model.encoder.layers.{i}.blocks.{j}.layernorm_before.bias')    )
               rename_keys.append((F'backbone.0.body.layers.{i}.blocks.{j}.attn.relative_position_bias_table', F'model.backbone.model.encoder.layers.{i}.blocks.{j}.attention.self.relative_position_bias_table')    )
               rename_keys.append((F'backbone.0.body.layers.{i}.blocks.{j}.attn.relative_position_index', F'model.backbone.model.encoder.layers.{i}.blocks.{j}.attention.self.relative_position_index')    )
               rename_keys.append((F'backbone.0.body.layers.{i}.blocks.{j}.attn.proj.weight', F'model.backbone.model.encoder.layers.{i}.blocks.{j}.attention.output.dense.weight')    )
               rename_keys.append((F'backbone.0.body.layers.{i}.blocks.{j}.attn.proj.bias', F'model.backbone.model.encoder.layers.{i}.blocks.{j}.attention.output.dense.bias')    )
               rename_keys.append((F'backbone.0.body.layers.{i}.blocks.{j}.norm2.weight', F'model.backbone.model.encoder.layers.{i}.blocks.{j}.layernorm_after.weight')    )
               rename_keys.append((F'backbone.0.body.layers.{i}.blocks.{j}.norm2.bias', F'model.backbone.model.encoder.layers.{i}.blocks.{j}.layernorm_after.bias')    )
               rename_keys.append((F'backbone.0.body.layers.{i}.blocks.{j}.mlp.fc1.weight', F'model.backbone.model.encoder.layers.{i}.blocks.{j}.intermediate.dense.weight')    )
               rename_keys.append((F'backbone.0.body.layers.{i}.blocks.{j}.mlp.fc1.bias', F'model.backbone.model.encoder.layers.{i}.blocks.{j}.intermediate.dense.bias')    )
               rename_keys.append((F'backbone.0.body.layers.{i}.blocks.{j}.mlp.fc2.weight', F'model.backbone.model.encoder.layers.{i}.blocks.{j}.output.dense.weight')    )
               rename_keys.append((F'backbone.0.body.layers.{i}.blocks.{j}.mlp.fc2.bias', F'model.backbone.model.encoder.layers.{i}.blocks.{j}.output.dense.bias')    )

          if i < 3:
               rename_keys.append((F'backbone.0.body.layers.{i}.downsample.reduction.weight', F'model.backbone.model.encoder.layers.{i}.downsample.reduction.weight')    )
               rename_keys.append((F'backbone.0.body.layers.{i}.downsample.norm.weight', F'model.backbone.model.encoder.layers.{i}.downsample.norm.weight')    )
               rename_keys.append((F'backbone.0.body.layers.{i}.downsample.norm.bias', F'model.backbone.model.encoder.layers.{i}.downsample.norm.bias')    )

     rename_keys.append(('backbone.0.body.norm1.weight', 'model.backbone.model.hidden_states_norms.stage2.weight')    )
     rename_keys.append(('backbone.0.body.norm1.bias', 'model.backbone.model.hidden_states_norms.stage2.bias')    )
     rename_keys.append(('backbone.0.body.norm2.weight', 'model.backbone.model.hidden_states_norms.stage3.weight')    )
     rename_keys.append(('backbone.0.body.norm2.bias', 'model.backbone.model.hidden_states_norms.stage3.bias')    )
     rename_keys.append(('backbone.0.body.norm3.weight', 'model.backbone.model.hidden_states_norms.stage4.weight')    )
     rename_keys.append(('backbone.0.body.norm3.bias', 'model.backbone.model.hidden_states_norms.stage4.bias')    )

     # transformer encoder
     for i in range(config.encoder_layers    ):
          rename_keys.append((F'transformer.encoder.layers.{i}.self_attn.sampling_offsets.weight', F'model.encoder.layers.{i}.self_attn.sampling_offsets.weight')    )
          rename_keys.append((F'transformer.encoder.layers.{i}.self_attn.sampling_offsets.bias', F'model.encoder.layers.{i}.self_attn.sampling_offsets.bias')    )
          rename_keys.append((F'transformer.encoder.layers.{i}.self_attn.attention_weights.weight', F'model.encoder.layers.{i}.self_attn.attention_weights.weight')    )
          rename_keys.append((F'transformer.encoder.layers.{i}.self_attn.attention_weights.bias', F'model.encoder.layers.{i}.self_attn.attention_weights.bias')    )
          rename_keys.append((F'transformer.encoder.layers.{i}.self_attn.value_proj.weight', F'model.encoder.layers.{i}.self_attn.value_proj.weight')    )
          rename_keys.append((F'transformer.encoder.layers.{i}.self_attn.value_proj.bias', F'model.encoder.layers.{i}.self_attn.value_proj.bias')    )
          rename_keys.append((F'transformer.encoder.layers.{i}.self_attn.output_proj.weight', F'model.encoder.layers.{i}.self_attn.output_proj.weight')    )
          rename_keys.append((F'transformer.encoder.layers.{i}.self_attn.output_proj.bias', F'model.encoder.layers.{i}.self_attn.output_proj.bias')    )
          rename_keys.append((F'transformer.encoder.layers.{i}.norm1.weight', F'model.encoder.layers.{i}.self_attn_layer_norm.weight')    )
          rename_keys.append((F'transformer.encoder.layers.{i}.norm1.bias', F'model.encoder.layers.{i}.self_attn_layer_norm.bias')    )
          rename_keys.append((F'transformer.encoder.layers.{i}.linear1.weight', F'model.encoder.layers.{i}.fc1.weight')    )
          rename_keys.append((F'transformer.encoder.layers.{i}.linear1.bias', F'model.encoder.layers.{i}.fc1.bias')    )
          rename_keys.append((F'transformer.encoder.layers.{i}.linear2.weight', F'model.encoder.layers.{i}.fc2.weight')    )
          rename_keys.append((F'transformer.encoder.layers.{i}.linear2.bias', F'model.encoder.layers.{i}.fc2.bias')    )
          rename_keys.append((F'transformer.encoder.layers.{i}.norm2.weight', F'model.encoder.layers.{i}.final_layer_norm.weight')    )
          rename_keys.append((F'transformer.encoder.layers.{i}.norm2.bias', F'model.encoder.layers.{i}.final_layer_norm.bias')    )

     # transformer decoder
     for i in range(config.decoder_layers    ):
          rename_keys.append((F'transformer.decoder.layers.{i}.cross_attn.sampling_offsets.weight', F'model.decoder.layers.{i}.encoder_attn.sampling_offsets.weight')    )
          rename_keys.append((F'transformer.decoder.layers.{i}.cross_attn.sampling_offsets.bias', F'model.decoder.layers.{i}.encoder_attn.sampling_offsets.bias')    )
          rename_keys.append((F'transformer.decoder.layers.{i}.cross_attn.attention_weights.weight', F'model.decoder.layers.{i}.encoder_attn.attention_weights.weight')    )
          rename_keys.append((F'transformer.decoder.layers.{i}.cross_attn.attention_weights.bias', F'model.decoder.layers.{i}.encoder_attn.attention_weights.bias')    )
          rename_keys.append((F'transformer.decoder.layers.{i}.cross_attn.value_proj.weight', F'model.decoder.layers.{i}.encoder_attn.value_proj.weight')    )
          rename_keys.append((F'transformer.decoder.layers.{i}.cross_attn.value_proj.bias', F'model.decoder.layers.{i}.encoder_attn.value_proj.bias')    )
          rename_keys.append((F'transformer.decoder.layers.{i}.cross_attn.output_proj.weight', F'model.decoder.layers.{i}.encoder_attn.output_proj.weight')    )
          rename_keys.append((F'transformer.decoder.layers.{i}.cross_attn.output_proj.bias', F'model.decoder.layers.{i}.encoder_attn.output_proj.bias')    )
          rename_keys.append((F'transformer.decoder.layers.{i}.norm1.weight', F'model.decoder.layers.{i}.encoder_attn_layer_norm.weight')    )
          rename_keys.append((F'transformer.decoder.layers.{i}.norm1.bias', F'model.decoder.layers.{i}.encoder_attn_layer_norm.bias')    )
          rename_keys.append((F'transformer.decoder.layers.{i}.self_attn.out_proj.weight', F'model.decoder.layers.{i}.self_attn.out_proj.weight')    )
          rename_keys.append((F'transformer.decoder.layers.{i}.self_attn.out_proj.bias', F'model.decoder.layers.{i}.self_attn.out_proj.bias')    )
          rename_keys.append((F'transformer.decoder.layers.{i}.norm2.weight', F'model.decoder.layers.{i}.self_attn_layer_norm.weight')    )
          rename_keys.append((F'transformer.decoder.layers.{i}.norm2.bias', F'model.decoder.layers.{i}.self_attn_layer_norm.bias')    )
          rename_keys.append((F'transformer.decoder.layers.{i}.linear1.weight', F'model.decoder.layers.{i}.fc1.weight')    )
          rename_keys.append((F'transformer.decoder.layers.{i}.linear1.bias', F'model.decoder.layers.{i}.fc1.bias')    )
          rename_keys.append((F'transformer.decoder.layers.{i}.linear2.weight', F'model.decoder.layers.{i}.fc2.weight')    )
          rename_keys.append((F'transformer.decoder.layers.{i}.linear2.bias', F'model.decoder.layers.{i}.fc2.bias')    )
          rename_keys.append((F'transformer.decoder.layers.{i}.norm3.weight', F'model.decoder.layers.{i}.final_layer_norm.weight')    )
          rename_keys.append((F'transformer.decoder.layers.{i}.norm3.bias', F'model.decoder.layers.{i}.final_layer_norm.bias')    )

     # fmt: on

     return rename_keys

def 						_UpperCAmelCase    (      SCREAMING_SNAKE_CASE__      : List[str]      ,					SCREAMING_SNAKE_CASE__      : Tuple      ,					SCREAMING_SNAKE_CASE__      : List[Any]    ):
     __UpperCamelCase   							=dct.pop(SCREAMING_SNAKE_CASE__    )
     __UpperCamelCase   							=val

def 						_UpperCAmelCase    (      SCREAMING_SNAKE_CASE__      : int      ,					SCREAMING_SNAKE_CASE__      : str    ):
     __UpperCamelCase   							=[int(backbone_config.embed_dim * 2**i    ) for i in range(len(backbone_config.depths    )    )]
     for i in range(len(backbone_config.depths    )    ):
          __UpperCamelCase   							=num_features[i]
          for j in range(backbone_config.depths[i]    ):
               # fmt: off
               # read in weights + bias of input projection layer (in original implementation, this is a single matrix + bias)
               __UpperCamelCase   							=state_dict.pop(F'backbone.0.body.layers.{i}.blocks.{j}.attn.qkv.weight'    )
               __UpperCamelCase   							=state_dict.pop(F'backbone.0.body.layers.{i}.blocks.{j}.attn.qkv.bias'    )
               # next, add query, keys and values (in that order) to the state dict
               __UpperCamelCase   							=in_proj_weight[:dim, :]
               __UpperCamelCase   							=in_proj_bias[: dim]
               __UpperCamelCase   							=in_proj_weight[
                   dim : dim * 2, :
               ]
               __UpperCamelCase   							=in_proj_bias[
                   dim : dim * 2
               ]
               __UpperCamelCase   							=in_proj_weight[
                   -dim :, :
               ]
               __UpperCamelCase   							=in_proj_bias[-dim :]
               # fmt: on

def 						_UpperCAmelCase    (      SCREAMING_SNAKE_CASE__      : List[Any]      ,					SCREAMING_SNAKE_CASE__      : List[str]    ):
     # transformer decoder self-attention layers
     __UpperCamelCase   							=config.d_model
     for i in range(config.decoder_layers    ):
          # read in weights + bias of input projection layer of self-attention
          __UpperCamelCase   							=state_dict.pop(F'transformer.decoder.layers.{i}.self_attn.in_proj_weight'    )
          __UpperCamelCase   							=state_dict.pop(F'transformer.decoder.layers.{i}.self_attn.in_proj_bias'    )
          # next, add query, keys and values (in that order) to the state dict
          __UpperCamelCase   							=in_proj_weight[:hidden_size, :]
          __UpperCamelCase   							=in_proj_bias[:hidden_size]
          __UpperCamelCase   							=in_proj_weight[
              hidden_size : hidden_size * 2, :
          ]
          __UpperCamelCase   							=in_proj_bias[hidden_size : hidden_size * 2]
          __UpperCamelCase   							=in_proj_weight[-hidden_size:, :]
          __UpperCamelCase   							=in_proj_bias[-hidden_size:]

def 						_UpperCAmelCase    (      ):
     __UpperCamelCase   							='http://images.cocodataset.org/val2017/000000039769.jpg'
     __UpperCamelCase   							=Image.open(requests.get(SCREAMING_SNAKE_CASE__      ,					stream=SCREAMING_SNAKE_CASE__    ).raw    )

     return im

@torch.no_grad()
def 						_UpperCAmelCase    (      SCREAMING_SNAKE_CASE__      : Tuple      ,					SCREAMING_SNAKE_CASE__      : List[Any]      ,					SCREAMING_SNAKE_CASE__      : Tuple    ):
     __UpperCamelCase   							=get_deta_config(SCREAMING_SNAKE_CASE__    )

     # load original state dict
     if model_name == "deta-swin-large":
          __UpperCamelCase   							=hf_hub_download(repo_id='nielsr/deta-checkpoints'      ,					filename='adet_swin_ft.pth'    )
     elif model_name == "deta-swin-large-o365":
          __UpperCamelCase   							=hf_hub_download(repo_id='jozhang97/deta-swin-l-o365'      ,					filename='deta_swin_pt_o365.pth'    )
     else:
          raise ValueError(F'Model name {model_name} not supported'    )

     __UpperCamelCase   							=torch.load(SCREAMING_SNAKE_CASE__      ,					map_location='cpu'    )['model']

     # original state dict
     for name, param in state_dict.items():
          print(SCREAMING_SNAKE_CASE__      ,					param.shape    )

     # rename keys
     __UpperCamelCase   							=create_rename_keys(SCREAMING_SNAKE_CASE__    )
     for src, dest in rename_keys:
          rename_key(SCREAMING_SNAKE_CASE__      ,					SCREAMING_SNAKE_CASE__      ,					SCREAMING_SNAKE_CASE__    )
     read_in_swin_q_k_v(SCREAMING_SNAKE_CASE__      ,					config.backbone_config    )
     read_in_decoder_q_k_v(SCREAMING_SNAKE_CASE__      ,					SCREAMING_SNAKE_CASE__    )

     # fix some prefixes
     for key in state_dict.copy().keys():
          if "transformer.decoder.class_embed" in key or "transformer.decoder.bbox_embed" in key:
               __UpperCamelCase   							=state_dict.pop(SCREAMING_SNAKE_CASE__    )
               __UpperCamelCase   							=val
          if "input_proj" in key:
               __UpperCamelCase   							=state_dict.pop(SCREAMING_SNAKE_CASE__    )
               __UpperCamelCase   							=val
          if "level_embed" in key or "pos_trans" in key or "pix_trans" in key or "enc_output" in key:
               __UpperCamelCase   							=state_dict.pop(SCREAMING_SNAKE_CASE__    )
               __UpperCamelCase   							=val

    # finally, create HuggingFace model and load state dict
     __UpperCamelCase   							=DetaForObjectDetection(SCREAMING_SNAKE_CASE__    )
     model.load_state_dict(SCREAMING_SNAKE_CASE__    )
     model.eval()

     __UpperCamelCase   							='cuda' if torch.cuda.is_available() else 'cpu'
     model.to(SCREAMING_SNAKE_CASE__    )

     # load image processor
     __UpperCamelCase   							=DetaImageProcessor(format='coco_detection'    )

     # verify our conversion on image
     __UpperCamelCase   							=prepare_img()
     __UpperCamelCase   							=processor(images=SCREAMING_SNAKE_CASE__      ,					return_tensors='pt'    )
     __UpperCamelCase   							=encoding['pixel_values']
     __UpperCamelCase   							=model(pixel_values.to(SCREAMING_SNAKE_CASE__    )    )

     # verify logits
     print('Logits:'      ,					outputs.logits[0, :3, :3]    )
     print('Boxes:'      ,					outputs.pred_boxes[0, :3, :3]    )
     if model_name == "deta-swin-large":
          __UpperCamelCase   							=torch.tensor(
              [[-7.6308, -2.8485, -5.3737], [-7.2037, -4.5505, -4.8027], [-7.2943, -4.2611, -4.6617]]    )
          __UpperCamelCase   							=torch.tensor([[0.4987, 0.4969, 0.9999], [0.2549, 0.5498, 0.4805], [0.5498, 0.2757, 0.0569]]    )
     elif model_name == "deta-swin-large-o365":
          __UpperCamelCase   							=torch.tensor(
              [[-8.0122, -3.5720, -4.9717], [-8.1547, -3.6886, -4.6389], [-7.6610, -3.6194, -5.0134]]    )
          __UpperCamelCase   							=torch.tensor([[0.2523, 0.5549, 0.4881], [0.7715, 0.4149, 0.4601], [0.5503, 0.2753, 0.0575]]    )
     assert torch.allclose(outputs.logits[0, :3, :3]      ,					expected_logits.to(SCREAMING_SNAKE_CASE__    )      ,					atol=1E-4    )
     assert torch.allclose(outputs.pred_boxes[0, :3, :3]      ,					expected_boxes.to(SCREAMING_SNAKE_CASE__    )      ,					atol=1E-4    )
     print('Everything ok!'    )

     if pytorch_dump_folder_path:
          # Save model and processor
          logger.info(F'Saving PyTorch model and processor to {pytorch_dump_folder_path}...'    )
          Path(SCREAMING_SNAKE_CASE__    ).mkdir(exist_ok=SCREAMING_SNAKE_CASE__    )
          model.save_pretrained(SCREAMING_SNAKE_CASE__    )
          processor.save_pretrained(SCREAMING_SNAKE_CASE__    )

     # Push to hub
     if push_to_hub:
          print('Pushing model and processor to hub...'    )
          model.push_to_hub(F'jozhang97/{model_name}'    )
          processor.push_to_hub(F'jozhang97/{model_name}'    )


if __name__ == "__main__":
   _A			=						argparse.ArgumentParser()

   parser.add_argument(
       '--model_name',
       type=str,
       default='deta-swin-large',
       choices=['deta-swin-large', 'deta-swin-large-o365'],
       help='Name of the model 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			=						parser.parse_args()
   convert_deta_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)





