


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				(    UpperCAmelCase   ,				unittest.TestCase      ):
		_a     :  Optional[Any]           =     DiTPipeline
		_a     :  int           =     CLASS_CONDITIONED_IMAGE_GENERATION_PARAMS
		_a     :  Tuple           =     PipelineTesterMixin.required_optional_params - {
		    'latents',
		    'num_images_per_prompt',
		    'callback',
		    'callback_steps',
		}
		_a     :  str           =     CLASS_CONDITIONED_IMAGE_GENERATION_BATCH_PARAMS
		_a     :  List[str]           =     False
		def      UpperCAmelCase__    (      self						)			->	int:




									"""simple docstring"""





									torch.manual_seed(0						)
									_UpperCAmelCase					=      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      ,   )
									_UpperCAmelCase					=      AutoencoderKL()
									_UpperCAmelCase					=      DDIMScheduler()
									_UpperCAmelCase					=      {'transformer': transformer.eval(), 'vae': vae.eval(), 'scheduler': scheduler}
									return components
		def      UpperCAmelCase__    (      self      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE=0						)			->	str:




									"""simple docstring"""





									if str(_SCREAMING_SNAKE_CASE						).startswith('mps'						):
																_UpperCAmelCase					=      torch.manual_seed(_SCREAMING_SNAKE_CASE						)
									else:
																_UpperCAmelCase					=      torch.Generator(device=_SCREAMING_SNAKE_CASE						).manual_seed(_SCREAMING_SNAKE_CASE						)
									_UpperCAmelCase					=      {
									    'class_labels': [1],
									    'generator': generator,
									    'num_inference_steps': 2,
									    'output_type': 'numpy',
									}
									return inputs
		def      UpperCAmelCase__    (      self						)			->	int:




									"""simple docstring"""





									_UpperCAmelCase					=      'cpu'

									_UpperCAmelCase					=      self.get_dummy_components()
									_UpperCAmelCase					=      self.pipeline_class(**_SCREAMING_SNAKE_CASE						)
									pipe.to(_SCREAMING_SNAKE_CASE						)
									pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE						)

									_UpperCAmelCase					=      self.get_dummy_inputs(_SCREAMING_SNAKE_CASE						)
									_UpperCAmelCase					=      pipe(**_SCREAMING_SNAKE_CASE						).images
									_UpperCAmelCase					=      image[0, -3:, -3:, -1]

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




									"""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      UpperCAmelCase__    (      self						)			->	List[str]:




									"""simple docstring"""





									self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1e-3						)



@require_torch_gpu
@slow
class 		__a				(    unittest.TestCase      ):
		def      UpperCAmelCase__    (      self						)			->	Any:




									"""simple docstring"""





									super().tearDown()
									gc.collect()
									torch.cuda.empty_cache()
		def      UpperCAmelCase__    (      self						)			->	Any:




									"""simple docstring"""





									_UpperCAmelCase					=      torch.manual_seed(0						)

									_UpperCAmelCase					=      DiTPipeline.from_pretrained('facebook/DiT-XL-2-256'						)
									pipe.to('cuda'						)

									_UpperCAmelCase					=      ['vase', 'umbrella', 'white shark', 'white wolf']
									_UpperCAmelCase					=      pipe.get_label_ids(_SCREAMING_SNAKE_CASE						)

									_UpperCAmelCase					=      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						):
																_UpperCAmelCase					=      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      UpperCAmelCase__    (      self						)			->	int:




									"""simple docstring"""





									_UpperCAmelCase					=      DiTPipeline.from_pretrained('facebook/DiT-XL-2-512'						)
									_UpperCAmelCase					=      DPMSolverMultistepScheduler.from_config(pipe.scheduler.config						)
									pipe.to('cuda'						)

									_UpperCAmelCase					=      ['vase', 'umbrella']
									_UpperCAmelCase					=      pipe.get_label_ids(_SCREAMING_SNAKE_CASE						)

									_UpperCAmelCase					=      torch.manual_seed(0						)
									_UpperCAmelCase					=      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						):
																_UpperCAmelCase					=      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




import warnings
from contextlib import contextmanager

from ....processing_utils import ProcessorMixin

class 		__a				(    UpperCAmelCase      ):
		_a     :  Optional[int]           =     'MCTCTFeatureExtractor'
		_a     :  int           =     'AutoTokenizer'
		def __init__(      self      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE						)			->	List[str]:




									"""simple docstring"""





									super().__init__(_SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE						)
									_UpperCAmelCase					=      self.feature_extractor
									_UpperCAmelCase					=      False
		def __call__(      self      ,   *_SCREAMING_SNAKE_CASE      ,   **_SCREAMING_SNAKE_CASE						)			->	Tuple:




									"""simple docstring"""





									if self._in_target_context_manager:
																return self.current_processor(*_SCREAMING_SNAKE_CASE      ,   **_SCREAMING_SNAKE_CASE						)

									if "raw_speech" in kwargs:
																warnings.warn('Using `raw_speech` as a keyword argument is deprecated. Use `audio` instead.'						)
																_UpperCAmelCase					=      kwargs.pop('raw_speech'						)
									else:
																_UpperCAmelCase					=      kwargs.pop('audio'      ,   _SCREAMING_SNAKE_CASE						)
									_UpperCAmelCase					=      kwargs.pop('sampling_rate'      ,   _SCREAMING_SNAKE_CASE						)
									_UpperCAmelCase					=      kwargs.pop('text'      ,   _SCREAMING_SNAKE_CASE						)
									if len(_SCREAMING_SNAKE_CASE						) > 0:
																_UpperCAmelCase					=      args[0]
																_UpperCAmelCase					=      args[1:]

									if audio is None and text is None:
																raise ValueError('You need to specify either an `audio` or `text` input to process.'						)

									if audio is not None:
																_UpperCAmelCase					=      self.feature_extractor(_SCREAMING_SNAKE_CASE      ,   *_SCREAMING_SNAKE_CASE      ,   sampling_rate=_SCREAMING_SNAKE_CASE      ,   **_SCREAMING_SNAKE_CASE						)
									if text is not None:
																_UpperCAmelCase					=      self.tokenizer(_SCREAMING_SNAKE_CASE      ,   **_SCREAMING_SNAKE_CASE						)

									if text is None:
																return inputs
									elif audio is None:
																return encodings
									else:
																_UpperCAmelCase					=      encodings['input_ids']
																return inputs
		def      UpperCAmelCase__    (      self      ,   *_SCREAMING_SNAKE_CASE      ,   **_SCREAMING_SNAKE_CASE						)			->	Optional[int]:




									"""simple docstring"""





									return self.tokenizer.batch_decode(*_SCREAMING_SNAKE_CASE      ,   **_SCREAMING_SNAKE_CASE						)
		def      UpperCAmelCase__    (      self      ,   *_SCREAMING_SNAKE_CASE      ,   **_SCREAMING_SNAKE_CASE						)			->	str:




									"""simple docstring"""





									if self._in_target_context_manager:
																return self.current_processor.pad(*_SCREAMING_SNAKE_CASE      ,   **_SCREAMING_SNAKE_CASE						)

									_UpperCAmelCase					=      kwargs.pop('input_features'      ,   _SCREAMING_SNAKE_CASE						)
									_UpperCAmelCase					=      kwargs.pop('labels'      ,   _SCREAMING_SNAKE_CASE						)
									if len(_SCREAMING_SNAKE_CASE						) > 0:
																_UpperCAmelCase					=      args[0]
																_UpperCAmelCase					=      args[1:]

									if input_features is not None:
																_UpperCAmelCase					=      self.feature_extractor.pad(_SCREAMING_SNAKE_CASE      ,   *_SCREAMING_SNAKE_CASE      ,   **_SCREAMING_SNAKE_CASE						)
									if labels is not None:
																_UpperCAmelCase					=      self.tokenizer.pad(_SCREAMING_SNAKE_CASE      ,   **_SCREAMING_SNAKE_CASE						)

									if labels is None:
																return input_features
									elif input_features is None:
																return labels
									else:
																_UpperCAmelCase					=      labels['input_ids']
																return input_features
		def      UpperCAmelCase__    (      self      ,   *_SCREAMING_SNAKE_CASE      ,   **_SCREAMING_SNAKE_CASE						)			->	int:




									"""simple docstring"""





									return self.tokenizer.decode(*_SCREAMING_SNAKE_CASE      ,   **_SCREAMING_SNAKE_CASE						)






		@contextmanager
		def      UpperCAmelCase__    (      self						)			->	Optional[Any]:




									"""simple docstring"""





									warnings.warn(
									    '`as_target_processor` is deprecated and will be removed in v5 of Transformers. You can process your '
									    'labels by using the argument `text` of the regular `__call__` method (either in the same call as '
									    'your audio inputs, or in a separate call.'						)
									_UpperCAmelCase					=      True
									_UpperCAmelCase					=      self.tokenizer
									yield
									_UpperCAmelCase					=      self.feature_extractor
									_UpperCAmelCase					=      False
