

import gc
import unittest

import numpy as np
import torch

from diffusers import AutoencoderKL, DDIMScheduler, DiTPipeline, DPMSolverMultistepScheduler, TransformeraDModel
from diffusers.utils import is_xformers_available, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import (
    CLASS_CONDITIONED_IMAGE_GENERATION_BATCH_PARAMS,
    CLASS_CONDITIONED_IMAGE_GENERATION_PARAMS,
)
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()






class 	A_		(	_lowerCAmelCase      ,							unittest.TestCase    ):
							_UpperCAmelCase   :  Optional[int]										=				DiTPipeline
							_UpperCAmelCase   :  List[Any]										=				CLASS_CONDITIONED_IMAGE_GENERATION_PARAMS
							_UpperCAmelCase   :  Any										=				PipelineTesterMixin.required_optional_params - {
							    "latents",
							    "num_images_per_prompt",
							    "callback",
							    "callback_steps",
							}
							_UpperCAmelCase   :  Tuple										=				CLASS_CONDITIONED_IMAGE_GENERATION_BATCH_PARAMS
							_UpperCAmelCase   :  Any										=				False
							def    lowerCAmelCase					(		self		:      str):
								torch.manual_seed(0)
								__lowerCamelCase    :    List[Any] 	=			TransformeraDModel(
								    sample_size=1_6       ,num_layers=2       ,patch_size=4       ,attention_head_dim=8       ,num_attention_heads=2       ,in_channels=4       ,out_channels=8       ,attention_bias=_lowercase       ,activation_fn='gelu-approximate'       ,num_embeds_ada_norm=1_0_0_0       ,norm_type='ada_norm_zero'       ,norm_elementwise_affine=_lowercase       ,)
								__lowerCamelCase    :    Tuple 	=			AutoencoderKL()
								__lowerCamelCase    :    List[Any] 	=			DDIMScheduler()
								__lowerCamelCase    :    Tuple 	=			{'transformer': transformer.eval(), 'vae': vae.eval(), 'scheduler': scheduler}
								return components
							def    lowerCAmelCase					(		self		:      Optional[int]       ,SCREAMING_SNAKE_CASE__		:      List[Any]       ,SCREAMING_SNAKE_CASE__		:      str=0):
								if str(_lowercase).startswith('mps'):
									__lowerCamelCase    :    int 	=			torch.manual_seed(_lowercase)
								else:
									__lowerCamelCase    :    Any 	=			torch.Generator(device=_lowercase).manual_seed(_lowercase)
								__lowerCamelCase    :    Optional[int] 	=			{
								    'class_labels': [1],
								    'generator': generator,
								    'num_inference_steps': 2,
								    'output_type': 'numpy',
								}
								return inputs
							def    lowerCAmelCase					(		self		:      Any):
								__lowerCamelCase    :    Tuple 	=			'cpu'

								__lowerCamelCase    :    Dict 	=			self.get_dummy_components()
								__lowerCamelCase    :    str 	=			self.pipeline_class(**_lowercase)
								pipe.to(_lowercase)
								pipe.set_progress_bar_config(disable=_lowercase)

								__lowerCamelCase    :    Dict 	=			self.get_dummy_inputs(_lowercase)
								__lowerCamelCase    :    Any 	=			pipe(**_lowercase).images
								__lowerCamelCase    :    Tuple 	=			image[0, -3:, -3:, -1]

								self.assertEqual(image.shape       ,(1, 1_6, 1_6, 3))
								__lowerCamelCase    :    Tuple 	=			np.array([0.2946, 0.6601, 0.4329, 0.3296, 0.4144, 0.5319, 0.7273, 0.5013, 0.4457])
								__lowerCamelCase    :    Dict 	=			np.abs(image_slice.flatten() - expected_slice).max()
								self.assertLessEqual(_lowercase       ,1E-3)
							def    lowerCAmelCase					(		self		:      Any):
								self._test_inference_batch_single_identical(relax_max_difference=_lowercase       ,expected_max_diff=1E-3)


							@unittest.skipIf(
							    torch_device != 'cuda' or not is_xformers_available()       ,reason='XFormers attention is only available with CUDA and `xformers` installed'       ,)
							def    lowerCAmelCase					(		self		:      str):
								self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1E-3)






@require_torch_gpu
@slow
class 	A_		(	unittest.TestCase    ):
							def    lowerCAmelCase					(		self		:      Any):
								super().tearDown()
								gc.collect()
								torch.cuda.empty_cache()
							def    lowerCAmelCase					(		self		:      Any):
								__lowerCamelCase    :    List[str] 	=			torch.manual_seed(0)

								__lowerCamelCase    :    Union[str, Any] 	=			DiTPipeline.from_pretrained('facebook/DiT-XL-2-256')
								pipe.to('cuda')

								__lowerCamelCase    :    int 	=			['vase', 'umbrella', 'white shark', 'white wolf']
								__lowerCamelCase    :    List[Any] 	=			pipe.get_label_ids(_lowercase)

								__lowerCamelCase    :    Optional[int] 	=			pipe(_lowercase       ,generator=_lowercase       ,num_inference_steps=4_0       ,output_type='np').images

								for word, image in zip(_lowercase       ,_lowercase):
									__lowerCamelCase    :    int 	=			load_numpy(
									    F"https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/dit/{word}.npy")
									assert np.abs((expected_image - image).max()) < 1E-2


							def    lowerCAmelCase					(		self		:      Any):
								__lowerCamelCase    :    Union[str, Any] 	=			DiTPipeline.from_pretrained('facebook/DiT-XL-2-512')
								__lowerCamelCase    :    Tuple 	=			DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
								pipe.to('cuda')

								__lowerCamelCase    :    List[str] 	=			['vase', 'umbrella']
								__lowerCamelCase    :    str 	=			pipe.get_label_ids(_lowercase)

								__lowerCamelCase    :    Dict 	=			torch.manual_seed(0)
								__lowerCamelCase    :    Optional[Any] 	=			pipe(_lowercase       ,generator=_lowercase       ,num_inference_steps=2_5       ,output_type='np').images

								for word, image in zip(_lowercase       ,_lowercase):
									__lowerCamelCase    :    str 	=			load_numpy(
									    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
									    F"/dit/{word}_512.npy")

									assert np.abs((expected_image - image).max()) < 1E-1







from __future__ import annotations

import unittest

from transformers import is_tf_available, is_torch_available
from transformers.testing_utils import DUMMY_UNKNOWN_IDENTIFIER, SMALL_MODEL_IDENTIFIER, is_pt_tf_cross_test, slow


if is_tf_available():
					from transformers import (
					    AutoConfig,
					    BertConfig,
					    GPTaConfig,
					    TaConfig,
					    TFAutoModel,
					    TFAutoModelForCausalLM,
					    TFAutoModelForMaskedLM,
					    TFAutoModelForPreTraining,
					    TFAutoModelForQuestionAnswering,
					    TFAutoModelForSeqaSeqLM,
					    TFAutoModelForSequenceClassification,
					    TFAutoModelWithLMHead,
					    TFBertForMaskedLM,
					    TFBertForPreTraining,
					    TFBertForQuestionAnswering,
					    TFBertForSequenceClassification,
					    TFBertModel,
					    TFGPTaLMHeadModel,
					    TFRobertaForMaskedLM,
					    TFTaForConditionalGeneration,
					)
					from transformers.models.bert.modeling_tf_bert import TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST
					from transformers.models.gpta.modeling_tf_gpta import TF_GPT2_PRETRAINED_MODEL_ARCHIVE_LIST
					from transformers.models.ta.modeling_tf_ta import TF_T5_PRETRAINED_MODEL_ARCHIVE_LIST

if is_torch_available():
					from transformers import (
					    AutoModel,
					    AutoModelForCausalLM,
					    AutoModelForMaskedLM,
					    AutoModelForPreTraining,
					    AutoModelForQuestionAnswering,
					    AutoModelForSeqaSeqLM,
					    AutoModelForSequenceClassification,
					    AutoModelWithLMHead,
					    BertForMaskedLM,
					    BertForPreTraining,
					    BertForQuestionAnswering,
					    BertForSequenceClassification,
					    BertModel,
					    GPTaLMHeadModel,
					    RobertaForMaskedLM,
					    TaForConditionalGeneration,
					)






@is_pt_tf_cross_test
class 	A_		(	unittest.TestCase    ):
							@slow
							def    lowerCAmelCase					(		self		:      List[str]):
								# for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
								for model_name in ["bert-base-uncased"]:
									__lowerCamelCase    :    Optional[int] 	=			AutoConfig.from_pretrained(SCREAMING_SNAKE_CASE__)
									self.assertIsNotNone(SCREAMING_SNAKE_CASE__)
									self.assertIsInstance(SCREAMING_SNAKE_CASE__       ,SCREAMING_SNAKE_CASE__)

									__lowerCamelCase    :    Tuple 	=			TFAutoModel.from_pretrained(SCREAMING_SNAKE_CASE__       ,from_pt=SCREAMING_SNAKE_CASE__)
									self.assertIsNotNone(SCREAMING_SNAKE_CASE__)
									self.assertIsInstance(SCREAMING_SNAKE_CASE__       ,SCREAMING_SNAKE_CASE__)

									__lowerCamelCase    :    str 	=			AutoModel.from_pretrained(SCREAMING_SNAKE_CASE__       ,from_tf=SCREAMING_SNAKE_CASE__)
									self.assertIsNotNone(SCREAMING_SNAKE_CASE__)
									self.assertIsInstance(SCREAMING_SNAKE_CASE__       ,SCREAMING_SNAKE_CASE__)
							@slow
							def    lowerCAmelCase					(		self		:      Optional[int]):
								# for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
								for model_name in ["bert-base-uncased"]:
									__lowerCamelCase    :    str 	=			AutoConfig.from_pretrained(SCREAMING_SNAKE_CASE__)
									self.assertIsNotNone(SCREAMING_SNAKE_CASE__)
									self.assertIsInstance(SCREAMING_SNAKE_CASE__       ,SCREAMING_SNAKE_CASE__)

									__lowerCamelCase    :    Tuple 	=			TFAutoModelForPreTraining.from_pretrained(SCREAMING_SNAKE_CASE__       ,from_pt=SCREAMING_SNAKE_CASE__)
									self.assertIsNotNone(SCREAMING_SNAKE_CASE__)
									self.assertIsInstance(SCREAMING_SNAKE_CASE__       ,SCREAMING_SNAKE_CASE__)

									__lowerCamelCase    :    str 	=			AutoModelForPreTraining.from_pretrained(SCREAMING_SNAKE_CASE__       ,from_tf=SCREAMING_SNAKE_CASE__)
									self.assertIsNotNone(SCREAMING_SNAKE_CASE__)
									self.assertIsInstance(SCREAMING_SNAKE_CASE__       ,SCREAMING_SNAKE_CASE__)
							@slow
							def    lowerCAmelCase					(		self		:      Optional[int]):
								for model_name in TF_GPT2_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
									__lowerCamelCase    :    Union[str, Any] 	=			AutoConfig.from_pretrained(SCREAMING_SNAKE_CASE__)
									self.assertIsNotNone(SCREAMING_SNAKE_CASE__)
									self.assertIsInstance(SCREAMING_SNAKE_CASE__       ,SCREAMING_SNAKE_CASE__)

									__lowerCamelCase    :    Union[str, Any] 	=			TFAutoModelForCausalLM.from_pretrained(SCREAMING_SNAKE_CASE__       ,from_pt=SCREAMING_SNAKE_CASE__)
									__lowerCamelCase	,			__lowerCamelCase    :    Union[str, Any] 	=			TFAutoModelForCausalLM.from_pretrained(
									    SCREAMING_SNAKE_CASE__       ,output_loading_info=SCREAMING_SNAKE_CASE__       ,from_pt=SCREAMING_SNAKE_CASE__)
									self.assertIsNotNone(SCREAMING_SNAKE_CASE__)
									self.assertIsInstance(SCREAMING_SNAKE_CASE__       ,SCREAMING_SNAKE_CASE__)

									__lowerCamelCase    :    List[Any] 	=			AutoModelForCausalLM.from_pretrained(SCREAMING_SNAKE_CASE__       ,from_tf=SCREAMING_SNAKE_CASE__)
									__lowerCamelCase	,			__lowerCamelCase    :    Any 	=			AutoModelForCausalLM.from_pretrained(
									    SCREAMING_SNAKE_CASE__       ,output_loading_info=SCREAMING_SNAKE_CASE__       ,from_tf=SCREAMING_SNAKE_CASE__)
									self.assertIsNotNone(SCREAMING_SNAKE_CASE__)
									self.assertIsInstance(SCREAMING_SNAKE_CASE__       ,SCREAMING_SNAKE_CASE__)
							@slow
							def    lowerCAmelCase					(		self		:      List[Any]):
								for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
									__lowerCamelCase    :    List[str] 	=			AutoConfig.from_pretrained(SCREAMING_SNAKE_CASE__)
									self.assertIsNotNone(SCREAMING_SNAKE_CASE__)
									self.assertIsInstance(SCREAMING_SNAKE_CASE__       ,SCREAMING_SNAKE_CASE__)

									__lowerCamelCase    :    List[str] 	=			TFAutoModelWithLMHead.from_pretrained(SCREAMING_SNAKE_CASE__       ,from_pt=SCREAMING_SNAKE_CASE__)
									self.assertIsNotNone(SCREAMING_SNAKE_CASE__)
									self.assertIsInstance(SCREAMING_SNAKE_CASE__       ,SCREAMING_SNAKE_CASE__)

									__lowerCamelCase    :    Optional[Any] 	=			AutoModelWithLMHead.from_pretrained(SCREAMING_SNAKE_CASE__       ,from_tf=SCREAMING_SNAKE_CASE__)
									self.assertIsNotNone(SCREAMING_SNAKE_CASE__)
									self.assertIsInstance(SCREAMING_SNAKE_CASE__       ,SCREAMING_SNAKE_CASE__)
							@slow
							def    lowerCAmelCase					(		self		:      Tuple):
								for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
									__lowerCamelCase    :    Optional[Any] 	=			AutoConfig.from_pretrained(SCREAMING_SNAKE_CASE__)
									self.assertIsNotNone(SCREAMING_SNAKE_CASE__)
									self.assertIsInstance(SCREAMING_SNAKE_CASE__       ,SCREAMING_SNAKE_CASE__)

									__lowerCamelCase    :    Optional[int] 	=			TFAutoModelForMaskedLM.from_pretrained(SCREAMING_SNAKE_CASE__       ,from_pt=SCREAMING_SNAKE_CASE__)
									__lowerCamelCase	,			__lowerCamelCase    :    Dict 	=			TFAutoModelForMaskedLM.from_pretrained(
									    SCREAMING_SNAKE_CASE__       ,output_loading_info=SCREAMING_SNAKE_CASE__       ,from_pt=SCREAMING_SNAKE_CASE__)
									self.assertIsNotNone(SCREAMING_SNAKE_CASE__)
									self.assertIsInstance(SCREAMING_SNAKE_CASE__       ,SCREAMING_SNAKE_CASE__)

									__lowerCamelCase    :    List[Any] 	=			AutoModelForMaskedLM.from_pretrained(SCREAMING_SNAKE_CASE__       ,from_tf=SCREAMING_SNAKE_CASE__)
									__lowerCamelCase	,			__lowerCamelCase    :    Optional[Any] 	=			AutoModelForMaskedLM.from_pretrained(
									    SCREAMING_SNAKE_CASE__       ,output_loading_info=SCREAMING_SNAKE_CASE__       ,from_tf=SCREAMING_SNAKE_CASE__)
									self.assertIsNotNone(SCREAMING_SNAKE_CASE__)
									self.assertIsInstance(SCREAMING_SNAKE_CASE__       ,SCREAMING_SNAKE_CASE__)
							@slow
							def    lowerCAmelCase					(		self		:      str):
								for model_name in TF_T5_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
									__lowerCamelCase    :    Dict 	=			AutoConfig.from_pretrained(SCREAMING_SNAKE_CASE__)
									self.assertIsNotNone(SCREAMING_SNAKE_CASE__)
									self.assertIsInstance(SCREAMING_SNAKE_CASE__       ,SCREAMING_SNAKE_CASE__)

									__lowerCamelCase    :    Optional[Any] 	=			TFAutoModelForSeqaSeqLM.from_pretrained(SCREAMING_SNAKE_CASE__       ,from_pt=SCREAMING_SNAKE_CASE__)
									__lowerCamelCase	,			__lowerCamelCase    :    List[str] 	=			TFAutoModelForSeqaSeqLM.from_pretrained(
									    SCREAMING_SNAKE_CASE__       ,output_loading_info=SCREAMING_SNAKE_CASE__       ,from_pt=SCREAMING_SNAKE_CASE__)
									self.assertIsNotNone(SCREAMING_SNAKE_CASE__)
									self.assertIsInstance(SCREAMING_SNAKE_CASE__       ,SCREAMING_SNAKE_CASE__)

									__lowerCamelCase    :    Any 	=			AutoModelForSeqaSeqLM.from_pretrained(SCREAMING_SNAKE_CASE__       ,from_tf=SCREAMING_SNAKE_CASE__)
									__lowerCamelCase	,			__lowerCamelCase    :    List[str] 	=			AutoModelForSeqaSeqLM.from_pretrained(
									    SCREAMING_SNAKE_CASE__       ,output_loading_info=SCREAMING_SNAKE_CASE__       ,from_tf=SCREAMING_SNAKE_CASE__)
									self.assertIsNotNone(SCREAMING_SNAKE_CASE__)
									self.assertIsInstance(SCREAMING_SNAKE_CASE__       ,SCREAMING_SNAKE_CASE__)
							@slow
							def    lowerCAmelCase					(		self		:      Optional[Any]):
								# for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
								for model_name in ["bert-base-uncased"]:
									__lowerCamelCase    :    List[Any] 	=			AutoConfig.from_pretrained(SCREAMING_SNAKE_CASE__)
									self.assertIsNotNone(SCREAMING_SNAKE_CASE__)
									self.assertIsInstance(SCREAMING_SNAKE_CASE__       ,SCREAMING_SNAKE_CASE__)

									__lowerCamelCase    :    Optional[int] 	=			TFAutoModelForSequenceClassification.from_pretrained(SCREAMING_SNAKE_CASE__       ,from_pt=SCREAMING_SNAKE_CASE__)
									self.assertIsNotNone(SCREAMING_SNAKE_CASE__)
									self.assertIsInstance(SCREAMING_SNAKE_CASE__       ,SCREAMING_SNAKE_CASE__)

									__lowerCamelCase    :    int 	=			AutoModelForSequenceClassification.from_pretrained(SCREAMING_SNAKE_CASE__       ,from_tf=SCREAMING_SNAKE_CASE__)
									self.assertIsNotNone(SCREAMING_SNAKE_CASE__)
									self.assertIsInstance(SCREAMING_SNAKE_CASE__       ,SCREAMING_SNAKE_CASE__)
							@slow
							def    lowerCAmelCase					(		self		:      Any):
								# for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
								for model_name in ["bert-base-uncased"]:
									__lowerCamelCase    :    List[Any] 	=			AutoConfig.from_pretrained(SCREAMING_SNAKE_CASE__)
									self.assertIsNotNone(SCREAMING_SNAKE_CASE__)
									self.assertIsInstance(SCREAMING_SNAKE_CASE__       ,SCREAMING_SNAKE_CASE__)

									__lowerCamelCase    :    str 	=			TFAutoModelForQuestionAnswering.from_pretrained(SCREAMING_SNAKE_CASE__       ,from_pt=SCREAMING_SNAKE_CASE__)
									self.assertIsNotNone(SCREAMING_SNAKE_CASE__)
									self.assertIsInstance(SCREAMING_SNAKE_CASE__       ,SCREAMING_SNAKE_CASE__)

									__lowerCamelCase    :    Optional[int] 	=			AutoModelForQuestionAnswering.from_pretrained(SCREAMING_SNAKE_CASE__       ,from_tf=SCREAMING_SNAKE_CASE__)
									self.assertIsNotNone(SCREAMING_SNAKE_CASE__)
									self.assertIsInstance(SCREAMING_SNAKE_CASE__       ,SCREAMING_SNAKE_CASE__)
							def    lowerCAmelCase					(		self		:      int):
								__lowerCamelCase    :    List[Any] 	=			TFAutoModelWithLMHead.from_pretrained(SCREAMING_SNAKE_CASE__       ,from_pt=SCREAMING_SNAKE_CASE__)
								self.assertIsInstance(SCREAMING_SNAKE_CASE__       ,SCREAMING_SNAKE_CASE__)
								self.assertEqual(model.num_parameters()       ,1_4_4_1_0)
								self.assertEqual(model.num_parameters(only_trainable=SCREAMING_SNAKE_CASE__)       ,1_4_4_1_0)

								__lowerCamelCase    :    Union[str, Any] 	=			AutoModelWithLMHead.from_pretrained(SCREAMING_SNAKE_CASE__       ,from_tf=SCREAMING_SNAKE_CASE__)
								self.assertIsInstance(SCREAMING_SNAKE_CASE__       ,SCREAMING_SNAKE_CASE__)
								self.assertEqual(model.num_parameters()       ,1_4_4_1_0)
								self.assertEqual(model.num_parameters(only_trainable=SCREAMING_SNAKE_CASE__)       ,1_4_4_1_0)


							def    lowerCAmelCase					(		self		:      Tuple):
								__lowerCamelCase    :    str 	=			TFAutoModelWithLMHead.from_pretrained(SCREAMING_SNAKE_CASE__       ,from_pt=SCREAMING_SNAKE_CASE__)
								self.assertIsInstance(SCREAMING_SNAKE_CASE__       ,SCREAMING_SNAKE_CASE__)
								self.assertEqual(model.num_parameters()       ,1_4_4_1_0)
								self.assertEqual(model.num_parameters(only_trainable=SCREAMING_SNAKE_CASE__)       ,1_4_4_1_0)

								__lowerCamelCase    :    Optional[int] 	=			AutoModelWithLMHead.from_pretrained(SCREAMING_SNAKE_CASE__       ,from_tf=SCREAMING_SNAKE_CASE__)
								self.assertIsInstance(SCREAMING_SNAKE_CASE__       ,SCREAMING_SNAKE_CASE__)
								self.assertEqual(model.num_parameters()       ,1_4_4_1_0)
								self.assertEqual(model.num_parameters(only_trainable=SCREAMING_SNAKE_CASE__)       ,1_4_4_1_0)




