




"""simple docstring"""




import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    StableDiffusionSAGPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()

class   UpperCAmelCase_    (			snake_case   ,					snake_case   ,					unittest.TestCase		):
		UpperCamelCase							=StableDiffusionSAGPipeline
		UpperCamelCase							=TEXT_TO_IMAGE_PARAMS
		UpperCamelCase							=TEXT_TO_IMAGE_BATCH_PARAMS
		UpperCamelCase							=TEXT_TO_IMAGE_IMAGE_PARAMS
		UpperCamelCase							=TEXT_TO_IMAGE_IMAGE_PARAMS
		UpperCamelCase							=False
		def 	_lowerCamelCase							(						self       )						->       Optional[int]:
					torch.manual_seed(0       )
					__lowercase			:  Any								=       UNetaDConditionModel(
					    block_out_channels=(32, 64)				,						layers_per_block=2				,						sample_size=32				,						in_channels=4				,						out_channels=4				,						down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')				,						up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''')				,						cross_attention_dim=32				,						)
					__lowercase			:  Union[str, Any]								=       DDIMScheduler(
					    beta_start=0.0_0_0_8_5				,						beta_end=0.0_1_2				,						beta_schedule='''scaled_linear'''				,						clip_sample=UpperCamelCase_				,						set_alpha_to_one=UpperCamelCase_				,						)
					torch.manual_seed(0       )
					__lowercase			:  str								=       AutoencoderKL(
					    block_out_channels=[32, 64]				,						in_channels=3				,						out_channels=3				,						down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']				,						up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']				,						latent_channels=4				,						)
					torch.manual_seed(0       )
					__lowercase			:  str								=       CLIPTextConfig(
					    bos_token_id=0				,						eos_token_id=2				,						hidden_size=32				,						intermediate_size=37				,						layer_norm_eps=1E-05				,						num_attention_heads=4				,						num_hidden_layers=5				,						pad_token_id=1				,						vocab_size=10_00				,						)
					__lowercase			:  List[Any]								=       CLIPTextModel(UpperCamelCase_       )
					__lowercase			:  List[Any]								=       CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''       )

					__lowercase			:  Dict								=       {
					    '''unet''': unet,
					    '''scheduler''': scheduler,
					    '''vae''': vae,
					    '''text_encoder''': text_encoder,
					    '''tokenizer''': tokenizer,
					    '''safety_checker''': None,
					    '''feature_extractor''': None,
					}
					return components
		def 	_lowerCamelCase							(						self				,						UpperCamelCase_				,						UpperCamelCase_=0       )						->       List[str]:
					if str(UpperCamelCase_       ).startswith('''mps'''       ):
								__lowercase			:  Any								=       torch.manual_seed(UpperCamelCase_       )
					else:
								__lowercase			:  Optional[Any]								=       torch.Generator(device=UpperCamelCase_       ).manual_seed(UpperCamelCase_       )
					__lowercase			:  List[Any]								=       {
					    '''prompt''': '''.''',
					    '''generator''': generator,
					    '''num_inference_steps''': 2,
					    '''guidance_scale''': 1.0,
					    '''sag_scale''': 1.0,
					    '''output_type''': '''numpy''',
					}
					return inputs
		def 	_lowerCamelCase							(						self       )						->       List[Any]:
					super().test_inference_batch_single_identical(expected_max_diff=3E-3       )


@slow
@require_torch_gpu
class   UpperCAmelCase_    (			unittest.TestCase		):
		def 	_lowerCamelCase							(						self       )						->       str:
					# clean up the VRAM after each test
					super().tearDown()
					gc.collect()
					torch.cuda.empty_cache()
		def 	_lowerCamelCase							(						self       )						->       Union[str, Any]:
					__lowercase			:  List[Any]								=       StableDiffusionSAGPipeline.from_pretrained('''CompVis/stable-diffusion-v1-4'''       )
					__lowercase			:  Any								=       sag_pipe.to(UpperCamelCase_       )
					sag_pipe.set_progress_bar_config(disable=UpperCamelCase_       )

					__lowercase			:  Optional[Any]								=       '''.'''
					__lowercase			:  Union[str, Any]								=       torch.manual_seed(0       )
					__lowercase			:  Union[str, Any]								=       sag_pipe(
					    [prompt]				,						generator=UpperCamelCase_				,						guidance_scale=7.5				,						sag_scale=1.0				,						num_inference_steps=20				,						output_type='''np'''       )

					__lowercase			:  Any								=       output.images

					__lowercase			:  Any								=       image[0, -3:, -3:, -1]

					assert image.shape == (1, 5_12, 5_12, 3)
					__lowercase			:  Optional[Any]								=       np.array([0.1_5_6_8, 0.1_7_3_8, 0.1_6_9_5, 0.1_6_9_3, 0.1_5_0_7, 0.1_7_0_5, 0.1_5_4_7, 0.1_7_5_1, 0.1_9_4_9]       )

					assert np.abs(image_slice.flatten() - expected_slice       ).max() < 5E-2
		def 	_lowerCamelCase							(						self       )						->       List[Any]:
					__lowercase			:  List[Any]								=       StableDiffusionSAGPipeline.from_pretrained('''stabilityai/stable-diffusion-2-1-base'''       )
					__lowercase			:  Optional[int]								=       sag_pipe.to(UpperCamelCase_       )
					sag_pipe.set_progress_bar_config(disable=UpperCamelCase_       )

					__lowercase			:  Dict								=       '''.'''
					__lowercase			:  Optional[Any]								=       torch.manual_seed(0       )
					__lowercase			:  Union[str, Any]								=       sag_pipe(
					    [prompt]				,						generator=UpperCamelCase_				,						guidance_scale=7.5				,						sag_scale=1.0				,						num_inference_steps=20				,						output_type='''np'''       )

					__lowercase			:  List[str]								=       output.images

					__lowercase			:  Tuple								=       image[0, -3:, -3:, -1]

					assert image.shape == (1, 5_12, 5_12, 3)
					__lowercase			:  Tuple								=       np.array([0.3_4_5_9, 0.2_8_7_6, 0.2_5_3_7, 0.3_0_0_2, 0.2_6_7_1, 0.2_1_6_0, 0.3_0_2_6, 0.2_2_6_2, 0.2_3_7_1]       )

					assert np.abs(image_slice.flatten() - expected_slice       ).max() < 5E-2
		def 	_lowerCamelCase							(						self       )						->       Any:
					__lowercase			:  List[str]								=       StableDiffusionSAGPipeline.from_pretrained('''stabilityai/stable-diffusion-2-1-base'''       )
					__lowercase			:  Any								=       sag_pipe.to(UpperCamelCase_       )
					sag_pipe.set_progress_bar_config(disable=UpperCamelCase_       )

					__lowercase			:  List[Any]								=       '''.'''
					__lowercase			:  Tuple								=       torch.manual_seed(0       )
					__lowercase			:  Any								=       sag_pipe(
					    [prompt]				,						width=7_68				,						height=5_12				,						generator=UpperCamelCase_				,						guidance_scale=7.5				,						sag_scale=1.0				,						num_inference_steps=20				,						output_type='''np'''				,						)

					__lowercase			:  int								=       output.images

					assert image.shape == (1, 5_12, 7_68, 3)







'''simple docstring'''
from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_sentencepiece_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


if is_sentencepiece_available():
							from ..ta.tokenization_ta import TaTokenizer
else:
							from ...utils.dummy_sentencepiece_objects import TaTokenizer

__SCREAMING_SNAKE_CASE					:				str									=		TaTokenizer

if is_tokenizers_available():
							from ..ta.tokenization_ta_fast import TaTokenizerFast
else:
							from ...utils.dummy_tokenizers_objects import TaTokenizerFast

__SCREAMING_SNAKE_CASE					:				Any									=		TaTokenizerFast

__SCREAMING_SNAKE_CASE					:				Tuple									=		{'''configuration_mt5''': ['''MT5Config''', '''MT5OnnxConfig''']}

try:
							if not is_torch_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							__SCREAMING_SNAKE_CASE					:				List[Any]									=		[
							    '''MT5EncoderModel''',
							    '''MT5ForConditionalGeneration''',
							    '''MT5ForQuestionAnswering''',
							    '''MT5Model''',
							    '''MT5PreTrainedModel''',
							    '''MT5Stack''',
							]

try:
							if not is_tf_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							__SCREAMING_SNAKE_CASE					:				str									=		['''TFMT5EncoderModel''', '''TFMT5ForConditionalGeneration''', '''TFMT5Model''']

try:
							if not is_flax_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							__SCREAMING_SNAKE_CASE					:				List[str]									=		['''FlaxMT5EncoderModel''', '''FlaxMT5ForConditionalGeneration''', '''FlaxMT5Model''']


if TYPE_CHECKING:
							from .configuration_mta import MTaConfig, MTaOnnxConfig

							try:
														if not is_torch_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .modeling_mta import (
														    MTaEncoderModel,
														    MTaForConditionalGeneration,
														    MTaForQuestionAnswering,
														    MTaModel,
														    MTaPreTrainedModel,
														    MTaStack,
														)

							try:
														if not is_tf_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .modeling_tf_mta import TFMTaEncoderModel, TFMTaForConditionalGeneration, TFMTaModel

							try:
														if not is_flax_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .modeling_flax_mta import FlaxMTaEncoderModel, FlaxMTaForConditionalGeneration, FlaxMTaModel

else:
							import sys

							__SCREAMING_SNAKE_CASE					:				Union[str, Any]									=		_LazyModule(
							    __name__,
							    globals()['''__file__'''],
							    _import_structure,
							    extra_objects={'''MT5Tokenizer''': MTaTokenizer, '''MT5TokenizerFast''': MTaTokenizerFast},
							    module_spec=__spec__,
							)





