
import gc
import unittest

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

from diffusers import AltDiffusionPipeline, AutoencoderKL, DDIMScheduler, PNDMScheduler, UNetaDConditionModel
from diffusers.pipelines.alt_diffusion.modeling_roberta_series import (
    RobertaSeriesConfig,
    RobertaSeriesModelWithTransformation,
)
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 PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()





class     lowerCamelCase      (       __lowerCamelCase ,		__lowerCamelCase ,		__lowerCamelCase ,		unittest.TestCase					):
	UpperCamelCase_  :					Union[str, Any]    		=	AltDiffusionPipeline
	UpperCamelCase_  :					Optional[Any]    		=	TEXT_TO_IMAGE_PARAMS
	UpperCamelCase_  :					List[Any]    		=	TEXT_TO_IMAGE_BATCH_PARAMS
	UpperCamelCase_  :					List[str]    		=	TEXT_TO_IMAGE_IMAGE_PARAMS
	UpperCamelCase_  :					Optional[Any]    		=	TEXT_TO_IMAGE_IMAGE_PARAMS







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







				"""simple docstring"""






				torch.manual_seed(0			)
				SCREAMING_SNAKE_CASE     						= UNetaDConditionModel(
				    block_out_channels=(3_2, 6_4)   ,     layers_per_block=2   ,     sample_size=3_2   ,     in_channels=4   ,     out_channels=4   ,     down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')   ,     up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''')   ,     cross_attention_dim=3_2   ,     )
				SCREAMING_SNAKE_CASE     						= 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			)
				SCREAMING_SNAKE_CASE     						= AutoencoderKL(
				    block_out_channels=[3_2, 6_4]   ,     in_channels=3   ,     out_channels=3   ,     down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']   ,     up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']   ,     latent_channels=4   ,     )

				# TODO: address the non-deterministic text encoder (fails for save-load tests)
				# torch.manual_seed(0)
				# text_encoder_config = RobertaSeriesConfig(
				#     hidden_size=32,
				#     project_dim=32,
				#     intermediate_size=37,
				#     layer_norm_eps=1e-05,
				#     num_attention_heads=4,
				#     num_hidden_layers=5,
				#     vocab_size=5002,
				# )
				# text_encoder = RobertaSeriesModelWithTransformation(text_encoder_config)

				torch.manual_seed(0			)
				SCREAMING_SNAKE_CASE     						= CLIPTextConfig(
				    bos_token_id=0   ,     eos_token_id=2   ,     hidden_size=3_2   ,     projection_dim=3_2   ,     intermediate_size=3_7   ,     layer_norm_eps=1e-05   ,     num_attention_heads=4   ,     num_hidden_layers=5   ,     pad_token_id=1   ,     vocab_size=5_0_0_2   ,     )
				SCREAMING_SNAKE_CASE     						= CLIPTextModel(lowercase			)

				SCREAMING_SNAKE_CASE     						= XLMRobertaTokenizer.from_pretrained('''hf-internal-testing/tiny-xlm-roberta'''			)
				SCREAMING_SNAKE_CASE     						= 7_7

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







	def  snake_case__     (     self      :Tuple   ,     lowercase      :Any   ,     lowercase      :Union[str, Any]=0			)		->      Tuple:







				"""simple docstring"""






				if str(lowercase			).startswith('''mps'''			):
							SCREAMING_SNAKE_CASE     						= torch.manual_seed(lowercase			)
				else:
							SCREAMING_SNAKE_CASE     						= torch.Generator(device=lowercase			).manual_seed(lowercase			)
				SCREAMING_SNAKE_CASE     						= {
				    '''prompt''': '''A painting of a squirrel eating a burger''',
				    '''generator''': generator,
				    '''num_inference_steps''': 2,
				    '''guidance_scale''': 6.0,
				    '''output_type''': '''numpy''',
				}
				return inputs







	def  snake_case__     (     self      :Tuple			)		->      str:







				"""simple docstring"""






				super().test_attention_slicing_forward_pass(expected_max_diff=3e-3			)







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







				"""simple docstring"""






				super().test_inference_batch_single_identical(expected_max_diff=3e-3			)







	def  snake_case__     (     self      :int			)		->      List[str]:







				"""simple docstring"""






				SCREAMING_SNAKE_CASE     						= '''cpu'''  # ensure determinism for the device-dependent torch.Generator

				SCREAMING_SNAKE_CASE     						= self.get_dummy_components()
				torch.manual_seed(0			)
				SCREAMING_SNAKE_CASE     						= RobertaSeriesConfig(
				    hidden_size=3_2   ,     project_dim=3_2   ,     intermediate_size=3_7   ,     layer_norm_eps=1e-05   ,     num_attention_heads=4   ,     num_hidden_layers=5   ,     vocab_size=5_0_0_2   ,     )
				# TODO: remove after fixing the non-deterministic text encoder
				SCREAMING_SNAKE_CASE     						= RobertaSeriesModelWithTransformation(lowercase			)
				SCREAMING_SNAKE_CASE     						= text_encoder

				SCREAMING_SNAKE_CASE     						= AltDiffusionPipeline(**lowercase			)
				SCREAMING_SNAKE_CASE     						= alt_pipe.to(lowercase			)
				alt_pipe.set_progress_bar_config(disable=lowercase			)

				SCREAMING_SNAKE_CASE     						= self.get_dummy_inputs(lowercase			)
				SCREAMING_SNAKE_CASE     						= '''A photo of an astronaut'''
				SCREAMING_SNAKE_CASE     						= alt_pipe(**lowercase			)
				SCREAMING_SNAKE_CASE     						= output.images
				SCREAMING_SNAKE_CASE     						= image[0, -3:, -3:, -1]

				assert image.shape == (1, 6_4, 6_4, 3)
				SCREAMING_SNAKE_CASE     						= np.array(
				    [0.5_74_81_62, 0.60_44_71_45, 0.48_82_12_17, 0.50_10_06_36, 0.5_43_11_85, 0.45_76_36_83, 0.49_65_76_96, 0.48_13_27_33, 0.47_57_30_93]			)

				assert np.abs(image_slice.flatten() - expected_slice			).max() < 1e-2







	def  snake_case__     (     self      :List[Any]			)		->      Dict:







				"""simple docstring"""






				SCREAMING_SNAKE_CASE     						= '''cpu'''  # ensure determinism for the device-dependent torch.Generator

				SCREAMING_SNAKE_CASE     						= self.get_dummy_components()
				SCREAMING_SNAKE_CASE     						= PNDMScheduler(skip_prk_steps=lowercase			)
				torch.manual_seed(0			)
				SCREAMING_SNAKE_CASE     						= RobertaSeriesConfig(
				    hidden_size=3_2   ,     project_dim=3_2   ,     intermediate_size=3_7   ,     layer_norm_eps=1e-05   ,     num_attention_heads=4   ,     num_hidden_layers=5   ,     vocab_size=5_0_0_2   ,     )
				# TODO: remove after fixing the non-deterministic text encoder
				SCREAMING_SNAKE_CASE     						= RobertaSeriesModelWithTransformation(lowercase			)
				SCREAMING_SNAKE_CASE     						= text_encoder
				SCREAMING_SNAKE_CASE     						= AltDiffusionPipeline(**lowercase			)
				SCREAMING_SNAKE_CASE     						= alt_pipe.to(lowercase			)
				alt_pipe.set_progress_bar_config(disable=lowercase			)

				SCREAMING_SNAKE_CASE     						= self.get_dummy_inputs(lowercase			)
				SCREAMING_SNAKE_CASE     						= alt_pipe(**lowercase			)
				SCREAMING_SNAKE_CASE     						= output.images
				SCREAMING_SNAKE_CASE     						= image[0, -3:, -3:, -1]

				assert image.shape == (1, 6_4, 6_4, 3)
				SCREAMING_SNAKE_CASE     						= np.array(
				    [0.51_60_50_93, 0.5_70_72_41, 0.47_36_55_07, 0.50_57_88_86, 0.5_63_38_77, 0.4_64_25_03, 0.5_18_20_81, 0.48_76_34_84, 0.49_08_42_37]			)

				assert np.abs(image_slice.flatten() - expected_slice			).max() < 1e-2





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







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







				"""simple docstring"""






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







	def  snake_case__     (     self      :List[Any]			)		->      Any:







				"""simple docstring"""






				SCREAMING_SNAKE_CASE     						= AltDiffusionPipeline.from_pretrained('''BAAI/AltDiffusion'''   ,     safety_checker=lowercase			)
				SCREAMING_SNAKE_CASE     						= alt_pipe.to(lowercase			)
				alt_pipe.set_progress_bar_config(disable=lowercase			)

				SCREAMING_SNAKE_CASE     						= '''A painting of a squirrel eating a burger'''
				SCREAMING_SNAKE_CASE     						= torch.manual_seed(0			)
				SCREAMING_SNAKE_CASE     						= alt_pipe([prompt]   ,     generator=lowercase   ,     guidance_scale=6.0   ,     num_inference_steps=2_0   ,     output_type='''np'''			)

				SCREAMING_SNAKE_CASE     						= output.images

				SCREAMING_SNAKE_CASE     						= image[0, -3:, -3:, -1]

				assert image.shape == (1, 5_1_2, 5_1_2, 3)
				SCREAMING_SNAKE_CASE     						= np.array([0.10_10, 0.08_00, 0.07_94, 0.08_85, 0.08_43, 0.07_62, 0.07_69, 0.07_29, 0.05_86]			)

				assert np.abs(image_slice.flatten() - expected_slice			).max() < 1e-2







	def  snake_case__     (     self      :int			)		->      Optional[int]:







				"""simple docstring"""






				SCREAMING_SNAKE_CASE     						= DDIMScheduler.from_pretrained('''BAAI/AltDiffusion'''   ,     subfolder='''scheduler'''			)

				SCREAMING_SNAKE_CASE     						= AltDiffusionPipeline.from_pretrained('''BAAI/AltDiffusion'''   ,     scheduler=lowercase   ,     safety_checker=lowercase			)
				SCREAMING_SNAKE_CASE     						= alt_pipe.to(lowercase			)
				alt_pipe.set_progress_bar_config(disable=lowercase			)

				SCREAMING_SNAKE_CASE     						= '''A painting of a squirrel eating a burger'''
				SCREAMING_SNAKE_CASE     						= torch.manual_seed(0			)

				SCREAMING_SNAKE_CASE     						= alt_pipe([prompt]   ,     generator=lowercase   ,     num_inference_steps=2   ,     output_type='''numpy'''			)
				SCREAMING_SNAKE_CASE     						= output.images

				SCREAMING_SNAKE_CASE     						= image[0, -3:, -3:, -1]

				assert image.shape == (1, 5_1_2, 5_1_2, 3)
				SCREAMING_SNAKE_CASE     						= np.array([0.40_19, 0.40_52, 0.38_10, 0.41_19, 0.39_16, 0.39_82, 0.46_51, 0.41_95, 0.53_23]			)

				assert np.abs(image_slice.flatten() - expected_slice			).max() < 1e-2

from ...configuration_utils import PretrainedConfig
from ...utils import logging


__lowerCAmelCase					   =       logging.get_logger(__name__)

__lowerCAmelCase					   =       {
    'alibaba-damo/mgp-str-base': 'https://huggingface.co/alibaba-damo/mgp-str-base/resolve/main/config.json',
}





class     lowerCamelCase      (       __lowerCamelCase					):
	UpperCamelCase_  :					Union[str, Any]    		=	'mgp-str'







	def __init__(     self      :Union[str, Any]   ,     lowercase      :Union[str, Any]=[3_2, 1_2_8]   ,     lowercase      :Tuple=4   ,     lowercase      :Optional[int]=3   ,     lowercase      :Union[str, Any]=2_7   ,     lowercase      :List[str]=3_8   ,     lowercase      :Optional[Any]=5_0_2_5_7   ,     lowercase      :str=3_0_5_2_2   ,     lowercase      :Dict=7_6_8   ,     lowercase      :List[Any]=1_2   ,     lowercase      :List[str]=1_2   ,     lowercase      :Optional[Any]=4.0   ,     lowercase      :Optional[Any]=True   ,     lowercase      :str=False   ,     lowercase      :Optional[Any]=1e-5   ,     lowercase      :List[str]=0.0   ,     lowercase      :List[Any]=0.0   ,     lowercase      :Optional[int]=0.0   ,     lowercase      :Dict=False   ,     lowercase      :Optional[Any]=0.02   ,     **lowercase      :Optional[Any]   ,     )		->      Dict:







				"""simple docstring"""






				super().__init__(**lowercase			)

				SCREAMING_SNAKE_CASE     						= image_size
				SCREAMING_SNAKE_CASE     						= patch_size
				SCREAMING_SNAKE_CASE     						= num_channels
				SCREAMING_SNAKE_CASE     						= max_token_length
				SCREAMING_SNAKE_CASE     						= num_character_labels
				SCREAMING_SNAKE_CASE     						= num_bpe_labels
				SCREAMING_SNAKE_CASE     						= num_wordpiece_labels
				SCREAMING_SNAKE_CASE     						= hidden_size
				SCREAMING_SNAKE_CASE     						= num_hidden_layers
				SCREAMING_SNAKE_CASE     						= num_attention_heads
				SCREAMING_SNAKE_CASE     						= mlp_ratio
				SCREAMING_SNAKE_CASE     						= distilled
				SCREAMING_SNAKE_CASE     						= layer_norm_eps
				SCREAMING_SNAKE_CASE     						= drop_rate
				SCREAMING_SNAKE_CASE     						= qkv_bias
				SCREAMING_SNAKE_CASE     						= attn_drop_rate
				SCREAMING_SNAKE_CASE     						= drop_path_rate
				SCREAMING_SNAKE_CASE     						= output_aa_attentions
				SCREAMING_SNAKE_CASE     						= initializer_range