





import inspect
import unittest

from transformers import is_torch_available, is_vision_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, slow, torch_device

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


if is_torch_available():
      import torch

      from transformers import MODEL_MAPPING, PoolFormerConfig, PoolFormerForImageClassification, PoolFormerModel
      from transformers.models.poolformer.modeling_poolformer import POOLFORMER_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
      from PIL import Image

      from transformers import PoolFormerImageProcessor

class 	__snake_case					(      _a ):






    '''simple docstring'''





    def 						A		( self    :		Dict				):

      """simple docstring"""





      __lowerCAmelCase								=					self.config_class(**self.inputs_dict				)
      self.parent.assertTrue(hasattr(snake_case      ,    '''hidden_sizes'''				)				)
      self.parent.assertTrue(hasattr(snake_case      ,    '''num_encoder_blocks'''				)				)


class 	__snake_case					:






    '''simple docstring'''





    def __init__( self    :		Optional[int]      ,    snake_case    :		Optional[int]      ,    snake_case    :		Tuple=13      ,    snake_case    :		List[str]=64      ,    snake_case    :		Tuple=3      ,    snake_case    :		List[str]=4      ,    snake_case    :		Dict=[2, 2, 2, 2]      ,    snake_case    :		Union[str, Any]=[8, 4, 2, 1]      ,    snake_case    :		Dict=[16, 32, 64, 128]      ,    snake_case    :		int=[1, 4, 8, 16]      ,    snake_case    :		str=False      ,    snake_case    :		Any=True      ,    snake_case    :		Tuple="gelu"      ,    snake_case    :		List[Any]=0.1      ,    snake_case    :		int=0.02      ,    snake_case    :		Optional[int]=3      ,    snake_case    :		List[str]=None      ,    ):

      """simple docstring"""





      __lowerCAmelCase								=					parent
      __lowerCAmelCase								=					batch_size
      __lowerCAmelCase								=					image_size
      __lowerCAmelCase								=					num_channels
      __lowerCAmelCase								=					num_encoder_blocks
      __lowerCAmelCase								=					sr_ratios
      __lowerCAmelCase								=					depths
      __lowerCAmelCase								=					hidden_sizes
      __lowerCAmelCase								=					downsampling_rates
      __lowerCAmelCase								=					is_training
      __lowerCAmelCase								=					use_labels
      __lowerCAmelCase								=					hidden_act
      __lowerCAmelCase								=					hidden_dropout_prob
      __lowerCAmelCase								=					initializer_range
      __lowerCAmelCase								=					num_labels
      __lowerCAmelCase								=					scope





    def 						A		( self    :		Any				):

      """simple docstring"""





      __lowerCAmelCase								=					floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]				)

      __lowerCAmelCase								=					None
      if self.use_labels:
        __lowerCAmelCase								=					ids_tensor([self.batch_size, self.image_size, self.image_size]      ,    self.num_labels				)

      __lowerCAmelCase								=					PoolFormerConfig(
          image_size=self.image_size      ,    num_channels=self.num_channels      ,    num_encoder_blocks=self.num_encoder_blocks      ,    depths=self.depths      ,    hidden_sizes=self.hidden_sizes      ,    hidden_act=self.hidden_act      ,    hidden_dropout_prob=self.hidden_dropout_prob      ,    initializer_range=self.initializer_range      ,    )

      return config, pixel_values, labels





    def 						A		( self    :		List[Any]      ,    snake_case    :		List[Any]      ,    snake_case    :		Optional[int]      ,    snake_case    :		Optional[Any]				):

      """simple docstring"""





      __lowerCAmelCase								=					PoolFormerModel(config=snake_case				)
      model.to(snake_case				)
      model.eval()
      __lowerCAmelCase								=					model(snake_case				)
      __lowerCAmelCase								=					__lowerCAmelCase								=					self.image_size // 32.0
      self.parent.assertEqual(
          result.last_hidden_state.shape      ,    (self.batch_size, self.hidden_sizes[-1], expected_height, expected_width)				)





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

      """simple docstring"""





      __lowerCAmelCase								=					self.prepare_config_and_inputs()
      __lowerCAmelCase						,    __lowerCAmelCase						,    __lowerCAmelCase								=					config_and_inputs
      __lowerCAmelCase								=					{'''pixel_values''': pixel_values}
      return config, inputs_dict


@require_torch
class 	__snake_case					(      _a							,					_a							,					unittest.TestCase ):






    '''simple docstring'''
    _lowerCAmelCase  =						(PoolFormerModel, PoolFormerForImageClassification) if is_torch_available() else ()
    _lowerCAmelCase  =						(
        {"""feature-extraction""": PoolFormerModel, """image-classification""": PoolFormerForImageClassification}
        if is_torch_available()
        else {}
    )

    _lowerCAmelCase  =						False
    _lowerCAmelCase  =						False
    _lowerCAmelCase  =						False
    _lowerCAmelCase  =						False
    _lowerCAmelCase  =						False





    def 						A		( self    :		int				):

      """simple docstring"""





      __lowerCAmelCase								=					PoolFormerModelTester(self				)
      __lowerCAmelCase								=					PoolFormerConfigTester(self      ,    config_class=snake_case				)





    def 						A		( self    :		Optional[Any]				):

      """simple docstring"""





      self.config_tester.run_common_tests()





    def 						A		( self    :		List[Any]				):

      """simple docstring"""





      __lowerCAmelCase								=					self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_model(*snake_case				)





    @unittest.skip('''PoolFormer does not use inputs_embeds'''				)
    def 						A		( self    :		Optional[int]				):

      """simple docstring"""





      pass





    @unittest.skip('''PoolFormer does not have get_input_embeddings method and get_output_embeddings methods'''				)
    def 						A		( self    :		int				):

      """simple docstring"""





      pass





    def 						A		( self    :		Tuple				):

      """simple docstring"""





      def check_hidden_states_output(snake_case    :		Union[str, Any]      ,    snake_case    :		Tuple      ,    snake_case    :		List[str]				):
        __lowerCAmelCase								=					model_class(snake_case				)
        model.to(snake_case				)
        model.eval()

        with torch.no_grad():
          __lowerCAmelCase								=					model(**self._prepare_for_class(snake_case      ,    snake_case				)				)

        __lowerCAmelCase								=					outputs.hidden_states

        __lowerCAmelCase								=					self.model_tester.num_encoder_blocks
        self.assertEqual(len(snake_case				)      ,    snake_case				)

        # verify the first hidden states (first block)
        self.assertListEqual(
            list(hidden_states[0].shape[-3:]				)      ,    [
                self.model_tester.hidden_sizes[0],
                self.model_tester.image_size // 4,
                self.model_tester.image_size // 4,
            ]      ,    )

      __lowerCAmelCase						,    __lowerCAmelCase								=					self.model_tester.prepare_config_and_inputs_for_common()

      for model_class in self.all_model_classes:
        __lowerCAmelCase								=					True
        check_hidden_states_output(snake_case      ,    snake_case      ,    snake_case				)

        # check that output_hidden_states also work using config
        del inputs_dict["output_hidden_states"]
        __lowerCAmelCase								=					True

        check_hidden_states_output(snake_case      ,    snake_case      ,    snake_case				)





    def 						A		( self    :		Dict				):

      """simple docstring"""





      if not self.model_tester.is_training:
        return

      __lowerCAmelCase						,    __lowerCAmelCase								=					self.model_tester.prepare_config_and_inputs_for_common()
      __lowerCAmelCase								=					True

      for model_class in self.all_model_classes:
        if model_class in get_values(snake_case				):
          continue
        __lowerCAmelCase								=					model_class(snake_case				)
        model.to(snake_case				)
        model.train()
        __lowerCAmelCase								=					self._prepare_for_class(snake_case      ,    snake_case      ,    return_labels=snake_case				)
        __lowerCAmelCase								=					model(**snake_case				).loss
        loss.backward()





    def 						A		( self    :		int				):

      """simple docstring"""





      __lowerCAmelCase						,    __lowerCAmelCase								=					self.model_tester.prepare_config_and_inputs_for_common()

      for model_class in self.all_model_classes:
        __lowerCAmelCase								=					model_class(snake_case				)
        __lowerCAmelCase								=					inspect.signature(model.forward				)
        # signature.parameters is an OrderedDict => so arg_names order is deterministic
        __lowerCAmelCase								=					[*signature.parameters.keys()]

        __lowerCAmelCase								=					['''pixel_values''']
        self.assertListEqual(arg_names[:1]      ,    snake_case				)





    @slow
    def 						A		( self    :		str				):

      """simple docstring"""





      for model_name in POOLFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        __lowerCAmelCase								=					PoolFormerModel.from_pretrained(snake_case				)
        self.assertIsNotNone(snake_case				)


def        __SCREAMING_SNAKE_CASE	(				):



  '''simple docstring'''






  __lowerCAmelCase								=					Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png''')
  return image




@require_torch
class 	__snake_case					(      unittest.TestCase ):






    '''simple docstring'''





    @slow
    def 						A		( self    :		List[Any]				):

      """simple docstring"""





      __lowerCAmelCase								=					PoolFormerImageProcessor()
      __lowerCAmelCase								=					PoolFormerForImageClassification.from_pretrained('''sail/poolformer_s12'''				).to(snake_case				)

      __lowerCAmelCase								=					image_processor(images=prepare_img()      ,    return_tensors='''pt'''				).to(snake_case				)

      # forward pass
      with torch.no_grad():
        __lowerCAmelCase								=					model(**snake_case				)

      # verify the logits
      __lowerCAmelCase								=					torch.Size((1, 1_000)				)
      self.assertEqual(outputs.logits.shape      ,    snake_case				)

      __lowerCAmelCase								=					torch.tensor([-0.6113, 0.1685, -0.0492]				).to(snake_case				)
      self.assertTrue(torch.allclose(outputs.logits[0, :3]      ,    snake_case      ,    atol=1e-4				)				)







import argparse
import os

import torch
from torchvision.datasets.utils import download_url

from diffusers import AutoencoderKL, DDIMScheduler, DiTPipeline, TransformeraDModel


__lowercase				    =      {512: """DiT-XL-2-512x512.pt""", 256: """DiT-XL-2-256x256.pt"""}


def        __SCREAMING_SNAKE_CASE	(				__lowercase:		Any):



  '''simple docstring'''






  __lowerCAmelCase								=					F'''pretrained_models/{model_name}'''
  if not os.path.isfile(__lowercase):
    os.makedirs('''pretrained_models'''  ,				exist_ok=__lowercase)
    __lowerCAmelCase								=					F'''https://dl.fbaipublicfiles.com/DiT/models/{model_name}'''
    download_url(__lowercase  ,				'''pretrained_models''')
  __lowerCAmelCase								=					torch.load(__lowercase  ,				map_location=lambda __lowercase  ,				__lowercase: storage)
  return model


def        __SCREAMING_SNAKE_CASE	(				__lowercase:		Tuple):



  '''simple docstring'''






  __lowerCAmelCase								=					download_model(pretrained_models[args.image_size])

  __lowerCAmelCase								=					state_dict['''x_embedder.proj.weight''']
  __lowerCAmelCase								=					state_dict['''x_embedder.proj.bias''']
  state_dict.pop('''x_embedder.proj.weight''')
  state_dict.pop('''x_embedder.proj.bias''')

  for depth in range(28):
    __lowerCAmelCase								=					state_dict[
        '''t_embedder.mlp.0.weight'''
    ]
    __lowerCAmelCase								=					state_dict[
        '''t_embedder.mlp.0.bias'''
    ]
    __lowerCAmelCase								=					state_dict[
        '''t_embedder.mlp.2.weight'''
    ]
    __lowerCAmelCase								=					state_dict[
        '''t_embedder.mlp.2.bias'''
    ]
    __lowerCAmelCase								=					state_dict[
        '''y_embedder.embedding_table.weight'''
    ]

    __lowerCAmelCase								=					state_dict[
        F'''blocks.{depth}.adaLN_modulation.1.weight'''
    ]
    __lowerCAmelCase								=					state_dict[
        F'''blocks.{depth}.adaLN_modulation.1.bias'''
    ]

    __lowerCAmelCase						,    __lowerCAmelCase						,    __lowerCAmelCase								=					torch.chunk(state_dict[F'''blocks.{depth}.attn.qkv.weight''']  ,				3  ,				dim=0)
    __lowerCAmelCase						,    __lowerCAmelCase						,    __lowerCAmelCase								=					torch.chunk(state_dict[F'''blocks.{depth}.attn.qkv.bias''']  ,				3  ,				dim=0)

    __lowerCAmelCase								=					q
    __lowerCAmelCase								=					q_bias
    __lowerCAmelCase								=					k
    __lowerCAmelCase								=					k_bias
    __lowerCAmelCase								=					v
    __lowerCAmelCase								=					v_bias

    __lowerCAmelCase								=					state_dict[
        F'''blocks.{depth}.attn.proj.weight'''
    ]
    __lowerCAmelCase								=					state_dict[F'''blocks.{depth}.attn.proj.bias''']

    __lowerCAmelCase								=					state_dict[F'''blocks.{depth}.mlp.fc1.weight''']
    __lowerCAmelCase								=					state_dict[F'''blocks.{depth}.mlp.fc1.bias''']
    __lowerCAmelCase								=					state_dict[F'''blocks.{depth}.mlp.fc2.weight''']
    __lowerCAmelCase								=					state_dict[F'''blocks.{depth}.mlp.fc2.bias''']

    state_dict.pop(F'''blocks.{depth}.attn.qkv.weight''')
    state_dict.pop(F'''blocks.{depth}.attn.qkv.bias''')
    state_dict.pop(F'''blocks.{depth}.attn.proj.weight''')
    state_dict.pop(F'''blocks.{depth}.attn.proj.bias''')
    state_dict.pop(F'''blocks.{depth}.mlp.fc1.weight''')
    state_dict.pop(F'''blocks.{depth}.mlp.fc1.bias''')
    state_dict.pop(F'''blocks.{depth}.mlp.fc2.weight''')
    state_dict.pop(F'''blocks.{depth}.mlp.fc2.bias''')
    state_dict.pop(F'''blocks.{depth}.adaLN_modulation.1.weight''')
    state_dict.pop(F'''blocks.{depth}.adaLN_modulation.1.bias''')

  state_dict.pop('''t_embedder.mlp.0.weight''')
  state_dict.pop('''t_embedder.mlp.0.bias''')
  state_dict.pop('''t_embedder.mlp.2.weight''')
  state_dict.pop('''t_embedder.mlp.2.bias''')
  state_dict.pop('''y_embedder.embedding_table.weight''')

  __lowerCAmelCase								=					state_dict['''final_layer.adaLN_modulation.1.weight''']
  __lowerCAmelCase								=					state_dict['''final_layer.adaLN_modulation.1.bias''']
  __lowerCAmelCase								=					state_dict['''final_layer.linear.weight''']
  __lowerCAmelCase								=					state_dict['''final_layer.linear.bias''']

  state_dict.pop('''final_layer.linear.weight''')
  state_dict.pop('''final_layer.linear.bias''')
  state_dict.pop('''final_layer.adaLN_modulation.1.weight''')
  state_dict.pop('''final_layer.adaLN_modulation.1.bias''')

  # DiT XL/2
  __lowerCAmelCase								=					TransformeraDModel(
      sample_size=args.image_size // 8  ,				num_layers=28  ,				attention_head_dim=72  ,				in_channels=4  ,				out_channels=8  ,				patch_size=2  ,				attention_bias=__lowercase  ,				num_attention_heads=16  ,				activation_fn='''gelu-approximate'''  ,				num_embeds_ada_norm=10_00  ,				norm_type='''ada_norm_zero'''  ,				norm_elementwise_affine=__lowercase  ,				)
  transformer.load_state_dict(__lowercase  ,				strict=__lowercase)

  __lowerCAmelCase								=					DDIMScheduler(
      num_train_timesteps=10_00  ,				beta_schedule='''linear'''  ,				prediction_type='''epsilon'''  ,				clip_sample=__lowercase  ,				)

  __lowerCAmelCase								=					AutoencoderKL.from_pretrained(args.vae_model)

  __lowerCAmelCase								=					DiTPipeline(transformer=__lowercase  ,				vae=__lowercase  ,				scheduler=__lowercase)

  if args.save:
    pipeline.save_pretrained(args.checkpoint_path)


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

      parser.add_argument(
          """--image_size""",
          default=256,
          type=int,
          required=False,
          help="""Image size of pretrained model, either 256 or 512.""",
      )
      parser.add_argument(
          """--vae_model""",
          default="""stabilityai/sd-vae-ft-ema""",
          type=str,
          required=False,
          help="""Path to pretrained VAE model, either stabilityai/sd-vae-ft-mse or stabilityai/sd-vae-ft-ema.""",
      )
      parser.add_argument(
          """--save""", default=True, type=bool, required=False, help="""Whether to save the converted pipeline or not."""
      )
      parser.add_argument(
          """--checkpoint_path""", default=None, type=str, required=True, help="""Path to the output pipeline."""
      )

      __lowercase				    =      parser.parse_args()
      main(args)
