



import unittest
from typing import Tuple

import torch

from diffusers.utils import floats_tensor, randn_tensor, torch_all_close, torch_device
from diffusers.utils.testing_utils import require_torch





@require_torch
class 		UpperCAmelCase						:





			@property
			def 	__lowerCAmelCase							(	self			):
									return self.get_dummy_input()





			@property
			def 	__lowerCAmelCase							(	self			):
									if self.block_type == "down":
															return (4, 32, 16, 16)
									elif self.block_type == "mid":
															return (4, 32, 32, 32)
									elif self.block_type == "up":
															return (4, 32, 64, 64)

									raise ValueError(F'''\'{self.block_type}\' is not a supported block_type. Set it to \'up\', \'mid\', or \'down\'.'''			)





			def 	__lowerCAmelCase							(	self       ,     _lowerCAmelCase=True       ,     _lowerCAmelCase=False       ,     _lowerCAmelCase=False       ,     _lowerCAmelCase=False       ,     ):
									_lowerCAmelCase 			=  4
									_lowerCAmelCase 			=  32
									_lowerCAmelCase 			=  (32, 32)

									_lowerCAmelCase 			=  torch.manual_seed(0			)
									_lowerCAmelCase 			=  torch.device(__lowerCamelCase			)
									_lowerCAmelCase 			=  (batch_size, num_channels) + sizes
									_lowerCAmelCase 			=  randn_tensor(__lowerCamelCase       ,     generator=__lowerCamelCase       ,     device=__lowerCamelCase			)
									_lowerCAmelCase 			=  {"hidden_states": hidden_states}

									if include_temb:
															_lowerCAmelCase 			=  128
															_lowerCAmelCase 			=  randn_tensor((batch_size, temb_channels)       ,     generator=__lowerCamelCase       ,     device=__lowerCamelCase			)

									if include_res_hidden_states_tuple:
															_lowerCAmelCase 			=  torch.manual_seed(1			)
															_lowerCAmelCase 			=  (randn_tensor(__lowerCamelCase       ,     generator=__lowerCamelCase       ,     device=__lowerCamelCase			),)

									if include_encoder_hidden_states:
															_lowerCAmelCase 			=  floats_tensor((batch_size, 32, 32)			).to(__lowerCamelCase			)

									if include_skip_sample:
															_lowerCAmelCase 			=  randn_tensor(((batch_size, 3) + sizes)       ,     generator=__lowerCamelCase       ,     device=__lowerCamelCase			)

									return dummy_input





			def 	__lowerCAmelCase							(	self			):
									_lowerCAmelCase 			=  {
									    "in_channels": 32,
									    "out_channels": 32,
									    "temb_channels": 128,
									}
									if self.block_type == "up":
															_lowerCAmelCase 			=  32

									if self.block_type == "mid":
															init_dict.pop('''out_channels'''			)

									_lowerCAmelCase 			=  self.dummy_input
									return init_dict, inputs_dict





			def 	__lowerCAmelCase							(	self       ,     _lowerCAmelCase			):
									_lowerCAmelCase 			=  self.prepare_init_args_and_inputs_for_common()
									_lowerCAmelCase 			=  self.block_class(**__lowerCamelCase			)
									unet_block.to(__lowerCamelCase			)
									unet_block.eval()

									with torch.no_grad():
															_lowerCAmelCase 			=  unet_block(**__lowerCamelCase			)

									if isinstance(__lowerCamelCase       ,     __lowerCamelCase			):
															_lowerCAmelCase 			=  output[0]

									self.assertEqual(output.shape       ,     self.output_shape			)

									_lowerCAmelCase 			=  output[0, -1, -3:, -3:]
									_lowerCAmelCase 			=  torch.tensor(__lowerCamelCase			).to(__lowerCamelCase			)
									assert torch_all_close(output_slice.flatten()       ,     __lowerCamelCase       ,     atol=5E-3			)





			@unittest.skipIf(torch_device == '''mps'''       ,     '''Training is not supported in mps'''			)
			def 	__lowerCAmelCase							(	self			):
									_lowerCAmelCase 			=  self.prepare_init_args_and_inputs_for_common()
									_lowerCAmelCase 			=  self.block_class(**__lowerCamelCase			)
									model.to(__lowerCamelCase			)
									model.train()
									_lowerCAmelCase 			=  model(**__lowerCamelCase			)

									if isinstance(__lowerCamelCase       ,     __lowerCamelCase			):
															_lowerCAmelCase 			=  output[0]

									_lowerCAmelCase 			=  torch.device(__lowerCamelCase			)
									_lowerCAmelCase 			=  randn_tensor(output.shape       ,     device=__lowerCamelCase			)
									_lowerCAmelCase 			=  torch.nn.functional.mse_loss(__lowerCamelCase       ,     __lowerCamelCase			)
									loss.backward()




from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_tf_available,
    is_torch_available,
    is_vision_available,
)


UpperCAmelCase_									=							{"configuration_deit": ["DEIT_PRETRAINED_CONFIG_ARCHIVE_MAP", "DeiTConfig", "DeiTOnnxConfig"]}

try:
    if not is_vision_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    UpperCAmelCase_									=							["DeiTFeatureExtractor"]
    UpperCAmelCase_									=							["DeiTImageProcessor"]

try:
    if not is_torch_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    UpperCAmelCase_									=							[
        "DEIT_PRETRAINED_MODEL_ARCHIVE_LIST",
        "DeiTForImageClassification",
        "DeiTForImageClassificationWithTeacher",
        "DeiTForMaskedImageModeling",
        "DeiTModel",
        "DeiTPreTrainedModel",
    ]

try:
    if not is_tf_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    UpperCAmelCase_									=							[
        "TF_DEIT_PRETRAINED_MODEL_ARCHIVE_LIST",
        "TFDeiTForImageClassification",
        "TFDeiTForImageClassificationWithTeacher",
        "TFDeiTForMaskedImageModeling",
        "TFDeiTModel",
        "TFDeiTPreTrainedModel",
    ]


if TYPE_CHECKING:
    from .configuration_deit import DEIT_PRETRAINED_CONFIG_ARCHIVE_MAP, DeiTConfig, DeiTOnnxConfig

    try:
        if not is_vision_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .feature_extraction_deit import DeiTFeatureExtractor
        from .image_processing_deit import DeiTImageProcessor

    try:
        if not is_torch_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .modeling_deit import (
            DEIT_PRETRAINED_MODEL_ARCHIVE_LIST,
            DeiTForImageClassification,
            DeiTForImageClassificationWithTeacher,
            DeiTForMaskedImageModeling,
            DeiTModel,
            DeiTPreTrainedModel,
        )

    try:
        if not is_tf_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .modeling_tf_deit import (
            TF_DEIT_PRETRAINED_MODEL_ARCHIVE_LIST,
            TFDeiTForImageClassification,
            TFDeiTForImageClassificationWithTeacher,
            TFDeiTForMaskedImageModeling,
            TFDeiTModel,
            TFDeiTPreTrainedModel,
        )


else:
    import sys

    UpperCAmelCase_									=							_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)