





'''simple docstring'''


import inspect
from typing import List, Optional, Tuple, Union

import numpy as np
import PIL
import torch
import torch.utils.checkpoint

from ...models import UNetaDModel, VQModel
from ...schedulers import (
    DDIMScheduler,
    DPMSolverMultistepScheduler,
    EulerAncestralDiscreteScheduler,
    EulerDiscreteScheduler,
    LMSDiscreteScheduler,
    PNDMScheduler,
)
from ...utils import PIL_INTERPOLATION, randn_tensor
from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput
def 							UpperCAmelCase_  (						__lowerCamelCase : Tuple			):
   lowercase_      :List[str]         =			image.size
   lowercase_      :Any         =			(x - x % 32 for x in (w, h))  # resize to integer multiple of 32
   lowercase_      :Any         =			image.resize((w, h)     ,resample=PIL_INTERPOLATION["lanczos"]			)
   lowercase_      :Optional[Any]         =			np.array(__lowerCamelCase			).astype(np.floataa			) / 255.0
   lowercase_      :Any         =			image[None].transpose(0     ,3     ,1     ,2			)
   lowercase_      :str         =			torch.from_numpy(__lowerCamelCase			)
   return 2.0 * image - 1.0


class 				a_       (    _lowerCAmelCase							):





     def __init__(			self		: Union[str, Any]     ,				lowercase		: VQModel     ,				lowercase		: UNetaDModel     ,				lowercase		: Union[
             DDIMScheduler,
             PNDMScheduler,
             LMSDiscreteScheduler,
             EulerDiscreteScheduler,
             EulerAncestralDiscreteScheduler,
             DPMSolverMultistepScheduler,
         ]     ,				):





        """simple docstring"""




        super().__init__()
        self.register_modules(vqvae=lowercase     ,				unet=lowercase     ,				scheduler=lowercase						)







     @torch.no_grad()
     def __call__(			self		: Any     ,				lowercase		: Union[torch.Tensor, PIL.Image.Image] = None     ,				lowercase		: Optional[int] = 1     ,				lowercase		: Optional[int] = 100     ,				lowercase		: Optional[float] = 0.0     ,				lowercase		: Optional[Union[torch.Generator, List[torch.Generator]]] = None     ,				lowercase		: Optional[str] = "pil"     ,				lowercase		: bool = True     ,				):





        """simple docstring"""




        if isinstance(lowercase     ,				PIL.Image.Image						):
           lowercase_      :Union[str, Any]         =			1
        elif isinstance(lowercase     ,				torch.Tensor						):
           lowercase_      :List[str]         =			image.shape[0]
        else:
           raise ValueError(F'`image` has to be of type `PIL.Image.Image` or `torch.Tensor` but is {type(lowercase						)}'						)

        if isinstance(lowercase     ,				PIL.Image.Image						):
           lowercase_      :str         =			preprocess(lowercase						)

        lowercase_      :Optional[int]         =			image.shape[-2:]

        # in_channels should be 6: 3 for latents, 3 for low resolution image
        lowercase_      :Optional[int]         =			(batch_size, self.unet.config.in_channels // 2, height, width)
        lowercase_      :Optional[int]         =			next(self.unet.parameters()						).dtype

        lowercase_      :Dict         =			randn_tensor(lowercase     ,				generator=lowercase     ,				device=self.device     ,				dtype=lowercase						)

        lowercase_      :List[Any]         =			image.to(device=self.device     ,				dtype=lowercase						)

        # set timesteps and move to the correct device
        self.scheduler.set_timesteps(lowercase     ,				device=self.device						)
        lowercase_      :Optional[int]         =			self.scheduler.timesteps

        # scale the initial noise by the standard deviation required by the scheduler
        lowercase_      :Tuple         =			latents * self.scheduler.init_noise_sigma

        # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature.
        # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
        # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
        # and should be between [0, 1]
        lowercase_      :Tuple         =			"eta" in set(inspect.signature(self.scheduler.step						).parameters.keys()						)
        lowercase_      :List[str]         =			{}
        if accepts_eta:
           lowercase_      :Dict         =			eta

        for t in self.progress_bar(lowercase						):
           # concat latents and low resolution image in the channel dimension.
           lowercase_      :int         =			torch.cat([latents, image]     ,				dim=1						)
           lowercase_      :Any         =			self.scheduler.scale_model_input(lowercase     ,				lowercase						)
           # predict the noise residual
           lowercase_      :List[Any]         =			self.unet(lowercase     ,				lowercase						).sample
           # compute the previous noisy sample x_t -> x_t-1
           lowercase_      :Tuple         =			self.scheduler.step(lowercase     ,				lowercase     ,				lowercase     ,				**lowercase						).prev_sample

        # decode the image latents with the VQVAE
        lowercase_      :Tuple         =			self.vqvae.decode(lowercase						).sample
        lowercase_      :int         =			torch.clamp(lowercase     ,				-1.0     ,				1.0						)
        lowercase_      :Any         =			image / 2 + 0.5
        lowercase_      :Tuple         =			image.cpu().permute(0     ,				2     ,				3     ,				1						).numpy()

        if output_type == "pil":
           lowercase_      :List[str]         =			self.numpy_to_pil(lowercase						)

        if not return_dict:
           return (image,)

        return ImagePipelineOutput(images=lowercase						)









'''simple docstring'''


from __future__ import annotations

import inspect
import unittest

import numpy as np

from transformers import DeiTConfig
from transformers.testing_utils import require_tf, require_vision, slow
from transformers.utils import cached_property, is_tf_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
 import tensorflow as tf

 from transformers import (
     TFDeiTForImageClassification,
     TFDeiTForImageClassificationWithTeacher,
     TFDeiTForMaskedImageModeling,
     TFDeiTModel,
 )
 from transformers.models.deit.modeling_tf_deit import TF_DEIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
 from PIL import Image

 from transformers import DeiTImageProcessor
class 				a_       :





     def __init__(			self		: int     ,				lowercase		: Optional[int]     ,				lowercase		: List[Any]=13     ,				lowercase		: Dict=30     ,				lowercase		: Dict=2     ,				lowercase		: Tuple=3     ,				lowercase		: Dict=True     ,				lowercase		: Dict=True     ,				lowercase		: Tuple=32     ,				lowercase		: List[Any]=2     ,				lowercase		: List[str]=4     ,				lowercase		: List[Any]=37     ,				lowercase		: Union[str, Any]="gelu"     ,				lowercase		: List[str]=0.1     ,				lowercase		: List[Any]=0.1     ,				lowercase		: List[str]=10     ,				lowercase		: Any=0.02     ,				lowercase		: Union[str, Any]=3     ,				lowercase		: Tuple=None     ,				lowercase		: List[str]=2     ,				):





        """simple docstring"""




        lowercase_      :Union[str, Any]         =			parent
        lowercase_      :Optional[int]         =			batch_size
        lowercase_      :Tuple         =			image_size
        lowercase_      :Any         =			patch_size
        lowercase_      :List[Any]         =			num_channels
        lowercase_      :Optional[Any]         =			is_training
        lowercase_      :str         =			use_labels
        lowercase_      :Any         =			hidden_size
        lowercase_      :Optional[int]         =			num_hidden_layers
        lowercase_      :List[Any]         =			num_attention_heads
        lowercase_      :str         =			intermediate_size
        lowercase_      :Optional[int]         =			hidden_act
        lowercase_      :List[Any]         =			hidden_dropout_prob
        lowercase_      :Optional[Any]         =			attention_probs_dropout_prob
        lowercase_      :Any         =			type_sequence_label_size
        lowercase_      :Dict         =			initializer_range
        lowercase_      :List[Any]         =			scope
        lowercase_      :List[str]         =			encoder_stride

        # in DeiT, the seq length equals the number of patches + 2 (we add 2 for the [CLS] and distilation tokens)
        lowercase_      :str         =			(image_size // patch_size) ** 2
        lowercase_      :Any         =			num_patches + 2





     def 							lowercase__				(			self		: List[Any]						):





        """simple docstring"""




        lowercase_      :List[str]         =			floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]						)

        lowercase_      :Optional[Any]         =			None
        if self.use_labels:
           lowercase_      :Optional[Any]         =			ids_tensor([self.batch_size]     ,				self.type_sequence_label_size						)

        lowercase_      :Any         =			self.get_config()

        return config, pixel_values, labels





     def 							lowercase__				(			self		: List[Any]						):





        """simple docstring"""




        return DeiTConfig(
            image_size=self.image_size     ,				patch_size=self.patch_size     ,				num_channels=self.num_channels     ,				hidden_size=self.hidden_size     ,				num_hidden_layers=self.num_hidden_layers     ,				num_attention_heads=self.num_attention_heads     ,				intermediate_size=self.intermediate_size     ,				hidden_act=self.hidden_act     ,				hidden_dropout_prob=self.hidden_dropout_prob     ,				attention_probs_dropout_prob=self.attention_probs_dropout_prob     ,				is_decoder=lowercase     ,				initializer_range=self.initializer_range     ,				encoder_stride=self.encoder_stride     ,				)





     def 							lowercase__				(			self		: Tuple     ,				lowercase		: int     ,				lowercase		: str     ,				lowercase		: Tuple						):





        """simple docstring"""




        lowercase_      :Optional[int]         =			TFDeiTModel(config=lowercase						)
        lowercase_      :str         =			model(lowercase						)
        self.parent.assertEqual(result.last_hidden_state.shape     ,				(self.batch_size, self.seq_length, self.hidden_size)						)





     def 							lowercase__				(			self		: str     ,				lowercase		: str     ,				lowercase		: Optional[Any]     ,				lowercase		: str						):





        """simple docstring"""




        lowercase_      :List[str]         =			TFDeiTForMaskedImageModeling(config=lowercase						)
        lowercase_      :Optional[int]         =			model(lowercase						)
        self.parent.assertEqual(
            result.reconstruction.shape     ,				(self.batch_size, self.num_channels, self.image_size, self.image_size)						)

        # test greyscale images
        lowercase_      :List[str]         =			1
        lowercase_      :List[str]         =			TFDeiTForMaskedImageModeling(lowercase						)

        lowercase_      :Union[str, Any]         =			floats_tensor([self.batch_size, 1, self.image_size, self.image_size]						)
        lowercase_      :List[str]         =			model(lowercase						)
        self.parent.assertEqual(result.reconstruction.shape     ,				(self.batch_size, 1, self.image_size, self.image_size)						)





     def 							lowercase__				(			self		: List[str]     ,				lowercase		: Any     ,				lowercase		: Union[str, Any]     ,				lowercase		: Union[str, Any]						):





        """simple docstring"""




        lowercase_      :Tuple         =			self.type_sequence_label_size
        lowercase_      :str         =			TFDeiTForImageClassification(lowercase						)
        lowercase_      :str         =			model(lowercase     ,				labels=lowercase						)
        self.parent.assertEqual(result.logits.shape     ,				(self.batch_size, self.type_sequence_label_size)						)

        # test greyscale images
        lowercase_      :Optional[int]         =			1
        lowercase_      :Optional[int]         =			TFDeiTForImageClassification(lowercase						)

        lowercase_      :List[Any]         =			floats_tensor([self.batch_size, 1, self.image_size, self.image_size]						)
        lowercase_      :int         =			model(lowercase     ,				labels=lowercase						)
        self.parent.assertEqual(result.logits.shape     ,				(self.batch_size, self.type_sequence_label_size)						)







     def 							lowercase__				(			self		: Dict						):





        """simple docstring"""




        lowercase_      :Tuple         =			self.prepare_config_and_inputs()
        lowercase_      ,							lowercase_      ,							lowercase_      :Tuple         =			config_and_inputs
        lowercase_      :Tuple         =			{"pixel_values": pixel_values}
        return config, inputs_dict







@require_tf
class 				a_       (    _lowerCAmelCase     ,						_lowerCAmelCase     ,						unittest.TestCase							):
     __A								=      (
         (
             TFDeiTModel,
             TFDeiTForImageClassification,
             TFDeiTForImageClassificationWithTeacher,
             TFDeiTForMaskedImageModeling,
         )
         if is_tf_available()
         else ()
     )
     __A								=      (
         {
             "feature-extraction": TFDeiTModel,
             "image-classification": (TFDeiTForImageClassification, TFDeiTForImageClassificationWithTeacher),
         }
         if is_tf_available()
         else {}
     )

     __A								=      False
     __A								=      False
     __A								=      False
     __A								=      False





     def 							lowercase__				(			self		: Optional[int]						):





        """simple docstring"""




        lowercase_      :List[Any]         =			TFDeiTModelTester(self						)
        lowercase_      :List[Any]         =			ConfigTester(self     ,				config_class=lowercase     ,				has_text_modality=lowercase     ,				hidden_size=37						)





     def 							lowercase__				(			self		: Tuple						):





        """simple docstring"""




        self.config_tester.run_common_tests()





     @unittest.skip(reason="DeiT does not use inputs_embeds"						)
     def 							lowercase__				(			self		: Optional[int]						):





        """simple docstring"""




        pass





     def 							lowercase__				(			self		: Tuple						):





        """simple docstring"""




        lowercase_      ,							lowercase_      :List[str]         =			self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
           lowercase_      :Optional[Any]         =			model_class(lowercase						)
           self.assertIsInstance(model.get_input_embeddings()     ,				(tf.keras.layers.Layer)						)
           lowercase_      :Tuple         =			model.get_output_embeddings()
           self.assertTrue(x is None or isinstance(lowercase     ,				tf.keras.layers.Dense						)						)





     def 							lowercase__				(			self		: Any						):





        """simple docstring"""




        lowercase_      ,							lowercase_      :List[str]         =			self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
           lowercase_      :Tuple         =			model_class(lowercase						)
           lowercase_      :Optional[int]         =			inspect.signature(model.call						)
           # signature.parameters is an OrderedDict => so arg_names order is deterministic
           lowercase_      :Tuple         =			[*signature.parameters.keys()]

           lowercase_      :Any         =			["pixel_values"]
           self.assertListEqual(arg_names[:1]     ,				lowercase						)





     def 							lowercase__				(			self		: Optional[Any]						):





        """simple docstring"""




        lowercase_      :Optional[Any]         =			self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*lowercase						)





     def 							lowercase__				(			self		: str						):





        """simple docstring"""




        lowercase_      :Optional[int]         =			self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_masked_image_modeling(*lowercase						)





     def 							lowercase__				(			self		: List[str]						):





        """simple docstring"""




        lowercase_      :Optional[int]         =			self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_image_classification(*lowercase						)





     def 							lowercase__				(			self		: List[Any]     ,				lowercase		: Optional[Any]     ,				lowercase		: List[str]     ,				lowercase		: Any=False						):





        """simple docstring"""




        lowercase_      :Dict         =			super()._prepare_for_class(lowercase     ,				lowercase     ,				return_labels=lowercase						)

        if return_labels:
           if "labels" in inputs_dict and "labels" not in inspect.signature(model_class.call						).parameters:
              del inputs_dict["labels"]

        return inputs_dict







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





        """simple docstring"""




        for model_name in TF_DEIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
           lowercase_      :List[Any]         =			TFDeiTModel.from_pretrained(lowercase						)
           self.assertIsNotNone(lowercase						)







def 							UpperCAmelCase_  (						):
   lowercase_      :List[str]         =			Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"			)
   return image


@require_tf
@require_vision
class 				a_       (    unittest.TestCase							):





     @cached_property
     def 							lowercase__				(			self		: Dict						):





        """simple docstring"""




        return (
            DeiTImageProcessor.from_pretrained("facebook/deit-base-distilled-patch16-224"						)
            if is_vision_available()
            else None
        )







     @slow
     def 							lowercase__				(			self		: Optional[Any]						):





        """simple docstring"""




        lowercase_      :Optional[int]         =			TFDeiTForImageClassificationWithTeacher.from_pretrained("facebook/deit-base-distilled-patch16-224"						)

        lowercase_      :Optional[int]         =			self.default_image_processor
        lowercase_      :Union[str, Any]         =			prepare_img()
        lowercase_      :Dict         =			image_processor(images=lowercase     ,				return_tensors="tf"						)

        # forward pass
        lowercase_      :Union[str, Any]         =			model(**lowercase						)

        # verify the logits
        lowercase_      :Union[str, Any]         =			tf.TensorShape((1, 1_000)						)
        self.assertEqual(outputs.logits.shape     ,				lowercase						)

        lowercase_      :Dict         =			tf.constant([-1.02_66, 0.19_12, -1.28_61]						)

        self.assertTrue(np.allclose(outputs.logits[0, :3]     ,				lowercase     ,				atol=1e-4						)						)


