



"""simple docstring"""






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__  (    __SCREAMING_SNAKE_CASE    ,						unittest.TestCase				):




	'''simple docstring'''





	SCREAMING_SNAKE_CASE        =      DiTPipeline
	SCREAMING_SNAKE_CASE        =      CLASS_CONDITIONED_IMAGE_GENERATION_PARAMS
	SCREAMING_SNAKE_CASE        =      PipelineTesterMixin.required_optional_params - {
	    'latents',
	    'num_images_per_prompt',
	    'callback',
	    'callback_steps',
	}
	SCREAMING_SNAKE_CASE        =      CLASS_CONDITIONED_IMAGE_GENERATION_BATCH_PARAMS
	SCREAMING_SNAKE_CASE        =      False

	def 				_SCREAMING_SNAKE_CASE  (	self:						Union[str, Any])							->			Dict:

				"""simple docstring"""

				torch.manual_seed(0)
				__lowerCAmelCase :    Optional[Any]				=			TransformeraDModel(
				    sample_size=16			,     num_layers=2			,     patch_size=4			,     attention_head_dim=8			,     num_attention_heads=2			,     in_channels=4			,     out_channels=8			,     attention_bias=_SCREAMING_SNAKE_CASE			,     activation_fn="gelu-approximate"			,     num_embeds_ada_norm=1000			,     norm_type="ada_norm_zero"			,     norm_elementwise_affine=_SCREAMING_SNAKE_CASE			,     )
				__lowerCAmelCase :    Optional[Any]				=			AutoencoderKL()
				__lowerCAmelCase :    List[Any]				=			DDIMScheduler()
				__lowerCAmelCase :    Dict				=			{"transformer": transformer.eval(), "vae": vae.eval(), "scheduler": scheduler}
				return components

	def 				_SCREAMING_SNAKE_CASE  (	self:						int			,     _SCREAMING_SNAKE_CASE:						List[str]			,     _SCREAMING_SNAKE_CASE:						Optional[int]=0)							->			Dict:

				"""simple docstring"""

				if str(_SCREAMING_SNAKE_CASE).startswith("mps"):
							__lowerCAmelCase :    str				=			torch.manual_seed(_SCREAMING_SNAKE_CASE)
				else:
							__lowerCAmelCase :    int				=			torch.Generator(device=_SCREAMING_SNAKE_CASE).manual_seed(_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase :    str				=			{
				    "class_labels": [1],
				    "generator": generator,
				    "num_inference_steps": 2,
				    "output_type": "numpy",
				}
				return inputs

	def 				_SCREAMING_SNAKE_CASE  (	self:						Optional[int])							->			Tuple:

				"""simple docstring"""

				__lowerCAmelCase :    Tuple				=			"cpu"

				__lowerCAmelCase :    Any				=			self.get_dummy_components()
				__lowerCAmelCase :    Any				=			self.pipeline_class(**_SCREAMING_SNAKE_CASE)
				pipe.to(_SCREAMING_SNAKE_CASE)
				pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE)

				__lowerCAmelCase :    List[Any]				=			self.get_dummy_inputs(_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase :    Tuple				=			pipe(**_SCREAMING_SNAKE_CASE).images
				__lowerCAmelCase :    Tuple				=			image[0, -3:, -3:, -1]

				self.assertEqual(image.shape			,     (1, 16, 16, 3))
				__lowerCAmelCase :    str				=			np.array([0.2946, 0.6601, 0.4329, 0.3296, 0.4144, 0.5319, 0.7273, 0.5013, 0.4457])
				__lowerCAmelCase :    Any				=			np.abs(image_slice.flatten() - expected_slice).max()
				self.assertLessEqual(_SCREAMING_SNAKE_CASE			,     1e-3)

	def 				_SCREAMING_SNAKE_CASE  (	self:						Optional[int])							->			str:

				"""simple docstring"""

				self._test_inference_batch_single_identical(relax_max_difference=_SCREAMING_SNAKE_CASE			,     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 				_SCREAMING_SNAKE_CASE  (	self:						str)							->			Dict:

				"""simple docstring"""

				self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1e-3)






@require_torch_gpu
@slow
class      A__  (    unittest.TestCase				):




	'''simple docstring'''





	def 				_SCREAMING_SNAKE_CASE  (	self:						Tuple)							->			Optional[Any]:

				"""simple docstring"""

				super().tearDown()
				gc.collect()
				torch.cuda.empty_cache()

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

				"""simple docstring"""

				__lowerCAmelCase :    List[Any]				=			torch.manual_seed(0)

				__lowerCAmelCase :    Any				=			DiTPipeline.from_pretrained("facebook/DiT-XL-2-256")
				pipe.to("cuda")

				__lowerCAmelCase :    Any				=			["vase", "umbrella", "white shark", "white wolf"]
				__lowerCAmelCase :    Optional[Any]				=			pipe.get_label_ids(_SCREAMING_SNAKE_CASE)

				__lowerCAmelCase :    Optional[Any]				=			pipe(_SCREAMING_SNAKE_CASE			,     generator=_SCREAMING_SNAKE_CASE			,     num_inference_steps=40			,     output_type="np").images

				for word, image in zip(_SCREAMING_SNAKE_CASE			,     _SCREAMING_SNAKE_CASE):
							__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 				_SCREAMING_SNAKE_CASE  (	self:						Dict)							->			List[str]:

				"""simple docstring"""

				__lowerCAmelCase :    Optional[Any]				=			DiTPipeline.from_pretrained("facebook/DiT-XL-2-512")
				__lowerCAmelCase :    Optional[int]				=			DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
				pipe.to("cuda")

				__lowerCAmelCase :    Tuple				=			["vase", "umbrella"]
				__lowerCAmelCase :    Dict				=			pipe.get_label_ids(_SCREAMING_SNAKE_CASE)

				__lowerCAmelCase :    List[str]				=			torch.manual_seed(0)
				__lowerCAmelCase :    Tuple				=			pipe(_SCREAMING_SNAKE_CASE			,     generator=_SCREAMING_SNAKE_CASE			,     num_inference_steps=25			,     output_type="np").images

				for word, image in zip(_SCREAMING_SNAKE_CASE			,     _SCREAMING_SNAKE_CASE):
							__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




"""simple docstring"""






import torch
import torch.nn as nn
from transformers.modeling_utils import ModuleUtilsMixin
from transformers.models.ta.modeling_ta import TaBlock, TaConfig, TaLayerNorm

from ...configuration_utils import ConfigMixin, register_to_config
from ...models import ModelMixin



class      A__  (    __SCREAMING_SNAKE_CASE    ,						__SCREAMING_SNAKE_CASE    ,						__SCREAMING_SNAKE_CASE				):




	'''simple docstring'''





	@register_to_config
	def __init__(	self:						List[str]			,     _SCREAMING_SNAKE_CASE:						int			,     _SCREAMING_SNAKE_CASE:						int			,     _SCREAMING_SNAKE_CASE:						int			,     _SCREAMING_SNAKE_CASE:						float			,     _SCREAMING_SNAKE_CASE:						int			,     _SCREAMING_SNAKE_CASE:						int			,     _SCREAMING_SNAKE_CASE:						int			,     _SCREAMING_SNAKE_CASE:						int			,     _SCREAMING_SNAKE_CASE:						str			,     _SCREAMING_SNAKE_CASE:						bool = False			,     )							->			Any:

				"""simple docstring"""

				super().__init__()

				__lowerCAmelCase :    str				=			nn.Embedding(_SCREAMING_SNAKE_CASE			,     _SCREAMING_SNAKE_CASE)

				__lowerCAmelCase :    Dict				=			nn.Embedding(_SCREAMING_SNAKE_CASE			,     _SCREAMING_SNAKE_CASE)
				__lowerCAmelCase :    str				=			False

				__lowerCAmelCase :    Optional[Any]				=			nn.Dropout(p=_SCREAMING_SNAKE_CASE)

				__lowerCAmelCase :    Dict				=			TaConfig(
				    vocab_size=_SCREAMING_SNAKE_CASE			,     d_model=_SCREAMING_SNAKE_CASE			,     num_heads=_SCREAMING_SNAKE_CASE			,     d_kv=_SCREAMING_SNAKE_CASE			,     d_ff=_SCREAMING_SNAKE_CASE			,     dropout_rate=_SCREAMING_SNAKE_CASE			,     feed_forward_proj=_SCREAMING_SNAKE_CASE			,     is_decoder=_SCREAMING_SNAKE_CASE			,     is_encoder_decoder=_SCREAMING_SNAKE_CASE			,     )

				__lowerCAmelCase :    str				=			nn.ModuleList()
				for lyr_num in range(_SCREAMING_SNAKE_CASE):
							__lowerCAmelCase :    int				=			TaBlock(_SCREAMING_SNAKE_CASE)
							self.encoders.append(_SCREAMING_SNAKE_CASE)

				__lowerCAmelCase :    Dict				=			TaLayerNorm(_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase :    Any				=			nn.Dropout(p=_SCREAMING_SNAKE_CASE)







	def 				_SCREAMING_SNAKE_CASE  (	self:						Any			,     _SCREAMING_SNAKE_CASE:						Tuple			,     _SCREAMING_SNAKE_CASE:						str)							->			List[str]:

				"""simple docstring"""

				__lowerCAmelCase :    Union[str, Any]				=			self.token_embedder(_SCREAMING_SNAKE_CASE)

				__lowerCAmelCase :    Optional[Any]				=			encoder_input_tokens.shape[1]
				__lowerCAmelCase :    List[Any]				=			torch.arange(_SCREAMING_SNAKE_CASE			,     device=encoder_input_tokens.device)
				x += self.position_encoding(_SCREAMING_SNAKE_CASE)

				__lowerCAmelCase :    List[Any]				=			self.dropout_pre(_SCREAMING_SNAKE_CASE)

				# inverted the attention mask
				__lowerCAmelCase :    List[Any]				=			encoder_input_tokens.size()
				__lowerCAmelCase :    Any				=			self.get_extended_attention_mask(_SCREAMING_SNAKE_CASE			,     _SCREAMING_SNAKE_CASE)

				for lyr in self.encoders:
							__lowerCAmelCase :    Union[str, Any]				=			lyr(_SCREAMING_SNAKE_CASE			,     _SCREAMING_SNAKE_CASE)[0]
				__lowerCAmelCase :    int				=			self.layer_norm(_SCREAMING_SNAKE_CASE)

				return self.dropout_post(_SCREAMING_SNAKE_CASE), encoder_inputs_mask