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 			lowerCamelCase_  (     UpperCAmelCase_		,			UpperCAmelCase_		,			unittest.TestCase				):






    '''simple docstring'''






    a__			:		Any      			=   StableDiffusionSAGPipeline
    a__			:		Optional[int]      			=   TEXT_TO_IMAGE_PARAMS
    a__			:		Optional[Any]      			=   TEXT_TO_IMAGE_BATCH_PARAMS
    a__			:		int      			=   TEXT_TO_IMAGE_IMAGE_PARAMS
    a__			:		int      			=   TEXT_TO_IMAGE_IMAGE_PARAMS
    a__			:		Optional[Any]      			=   False







    def     UpperCamelCase__     (			self) ->			Tuple:
        torch.manual_seed(0)
        __UpperCamelCase  :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				,				)
        __UpperCamelCase  :Any    						=	DDIMScheduler(
            beta_start=0.0_00_85				,				beta_end=0.0_12				,				beta_schedule='''scaled_linear'''				,				clip_sample=__lowercase				,				set_alpha_to_one=__lowercase				,				)
        torch.manual_seed(0)
        __UpperCamelCase  :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)
        __UpperCamelCase  :Optional[int]    						=	CLIPTextConfig(
            bos_token_id=0				,				eos_token_id=2				,				hidden_size=32				,				intermediate_size=37				,				layer_norm_eps=1E-0_5				,				num_attention_heads=4				,				num_hidden_layers=5				,				pad_token_id=1				,				vocab_size=1_000				,				)
        __UpperCamelCase  :List[Any]    						=	CLIPTextModel(__lowercase)
        __UpperCamelCase  :Union[str, Any]    						=	CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip''')

        __UpperCamelCase  :Dict    						=	{
            '''unet''': unet,
            '''scheduler''': scheduler,
            '''vae''': vae,
            '''text_encoder''': text_encoder,
            '''tokenizer''': tokenizer,
            '''safety_checker''': None,
            '''feature_extractor''': None,
        }
        return components







    def     UpperCamelCase__     (			self				,				__lowercase				,				__lowercase=0) ->			Tuple:
        if str(__lowercase).startswith('''mps'''):
            __UpperCamelCase  :Dict    						=	torch.manual_seed(__lowercase)
        else:
            __UpperCamelCase  :int    						=	torch.Generator(device=__lowercase).manual_seed(__lowercase)
        __UpperCamelCase  :Optional[int]    						=	{
            '''prompt''': '''.''',
            '''generator''': generator,
            '''num_inference_steps''': 2,
            '''guidance_scale''': 1.0,
            '''sag_scale''': 1.0,
            '''output_type''': '''numpy''',
        }
        return inputs







    def     UpperCamelCase__     (			self) ->			str:
        super().test_inference_batch_single_identical(expected_max_diff=3E-3)


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






    '''simple docstring'''






    def     UpperCamelCase__     (			self) ->			Union[str, Any]:
        # clean up the VRAM after each test
        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()







    def     UpperCamelCase__     (			self) ->			str:
        __UpperCamelCase  :Tuple    						=	StableDiffusionSAGPipeline.from_pretrained('''CompVis/stable-diffusion-v1-4''')
        __UpperCamelCase  :Tuple    						=	sag_pipe.to(__lowercase)
        sag_pipe.set_progress_bar_config(disable=__lowercase)

        __UpperCamelCase  :Tuple    						=	'''.'''
        __UpperCamelCase  :Dict    						=	torch.manual_seed(0)
        __UpperCamelCase  :Dict    						=	sag_pipe(
            [prompt]				,				generator=__lowercase				,				guidance_scale=7.5				,				sag_scale=1.0				,				num_inference_steps=20				,				output_type='''np''')

        __UpperCamelCase  :List[Any]    						=	output.images

        __UpperCamelCase  :Optional[Any]    						=	image[0, -3:, -3:, -1]

        assert image.shape == (1, 512, 512, 3)
        __UpperCamelCase  :Dict    						=	np.array([0.15_68, 0.17_38, 0.16_95, 0.16_93, 0.15_07, 0.17_05, 0.15_47, 0.17_51, 0.19_49])

        assert np.abs(image_slice.flatten() - expected_slice).max() < 5E-2







    def     UpperCamelCase__     (			self) ->			str:
        __UpperCamelCase  :Union[str, Any]    						=	StableDiffusionSAGPipeline.from_pretrained('''stabilityai/stable-diffusion-2-1-base''')
        __UpperCamelCase  :Union[str, Any]    						=	sag_pipe.to(__lowercase)
        sag_pipe.set_progress_bar_config(disable=__lowercase)

        __UpperCamelCase  :List[str]    						=	'''.'''
        __UpperCamelCase  :Dict    						=	torch.manual_seed(0)
        __UpperCamelCase  :Tuple    						=	sag_pipe(
            [prompt]				,				generator=__lowercase				,				guidance_scale=7.5				,				sag_scale=1.0				,				num_inference_steps=20				,				output_type='''np''')

        __UpperCamelCase  :List[str]    						=	output.images

        __UpperCamelCase  :List[str]    						=	image[0, -3:, -3:, -1]

        assert image.shape == (1, 512, 512, 3)
        __UpperCamelCase  :Union[str, Any]    						=	np.array([0.34_59, 0.28_76, 0.25_37, 0.30_02, 0.26_71, 0.21_60, 0.30_26, 0.22_62, 0.23_71])

        assert np.abs(image_slice.flatten() - expected_slice).max() < 5E-2







    def     UpperCamelCase__     (			self) ->			Any:
        __UpperCamelCase  :Any    						=	StableDiffusionSAGPipeline.from_pretrained('''stabilityai/stable-diffusion-2-1-base''')
        __UpperCamelCase  :Tuple    						=	sag_pipe.to(__lowercase)
        sag_pipe.set_progress_bar_config(disable=__lowercase)

        __UpperCamelCase  :List[str]    						=	'''.'''
        __UpperCamelCase  :str    						=	torch.manual_seed(0)
        __UpperCamelCase  :Any    						=	sag_pipe(
            [prompt]				,				width=768				,				height=512				,				generator=__lowercase				,				guidance_scale=7.5				,				sag_scale=1.0				,				num_inference_steps=20				,				output_type='''np'''				,				)

        __UpperCamelCase  :Dict    						=	output.images

        assert image.shape == (1, 512, 768, 3)






import unittest

from transformers import is_flax_available
from transformers.testing_utils import require_flax, require_sentencepiece, require_tokenizers, require_torch, slow


if is_flax_available():
    import optax
    from flax.training.common_utils import onehot

    from transformers import AutoTokenizer, FlaxMTaForConditionalGeneration
    from transformers.models.ta.modeling_flax_ta import shift_tokens_right


@require_torch
@require_sentencepiece
@require_tokenizers
@require_flax
class 			lowerCamelCase_  (     unittest.TestCase				):






    '''simple docstring'''






    @slow
    def     UpperCamelCase__     (			self) ->			List[str]:
        __UpperCamelCase  :List[str]    						=	FlaxMTaForConditionalGeneration.from_pretrained('''google/mt5-small''')
        __UpperCamelCase  :Dict    						=	AutoTokenizer.from_pretrained('''google/mt5-small''')

        __UpperCamelCase  :Optional[Any]    						=	tokenizer('''Hello there'''				,				return_tensors='''np''').input_ids
        __UpperCamelCase  :List[str]    						=	tokenizer('''Hi I am'''				,				return_tensors='''np''').input_ids

        __UpperCamelCase  :Optional[int]    						=	shift_tokens_right(__lowercase				,				model.config.pad_token_id				,				model.config.decoder_start_token_id)

        __UpperCamelCase  :Tuple    						=	model(__lowercase				,				decoder_input_ids=__lowercase).logits
        __UpperCamelCase  :Any    						=	optax.softmax_cross_entropy(__lowercase				,				onehot(__lowercase				,				logits.shape[-1])).mean()

        __UpperCamelCase  :str    						=	-(labels.shape[-1] * loss.item())

        __UpperCamelCase  :Optional[Any]    						=	-84.91_27
        self.assertTrue(abs(mtf_score - EXPECTED_SCORE) < 1E-4)





