




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

			"""simple docstring"""







			snake_case_  						=			AltDiffusionPipeline
			snake_case_  						=			TEXT_TO_IMAGE_PARAMS
			snake_case_  						=			TEXT_TO_IMAGE_BATCH_PARAMS
			snake_case_  						=			TEXT_TO_IMAGE_IMAGE_PARAMS
			snake_case_  						=			TEXT_TO_IMAGE_IMAGE_PARAMS


			def  lowercase_       ( self				)       -> Dict:



							'''simple docstring'''
							torch.manual_seed(0				)
							__lowerCamelCase									=       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     ,		)
							__lowerCamelCase									=       DDIMScheduler(
							    beta_start=0.0_00_85     ,		beta_end=0.0_12     ,		beta_schedule='scaled_linear'     ,		clip_sample=lowerCamelCase__     ,		set_alpha_to_one=lowerCamelCase__     ,		)
							torch.manual_seed(0				)
							__lowerCamelCase									=       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     ,		)

							# 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				)
							__lowerCamelCase									=       CLIPTextConfig(
							    bos_token_id=0     ,		eos_token_id=2     ,		hidden_size=32     ,		projection_dim=32     ,		intermediate_size=37     ,		layer_norm_eps=1e-05     ,		num_attention_heads=4     ,		num_hidden_layers=5     ,		pad_token_id=1     ,		vocab_size=5_002     ,		)
							__lowerCamelCase									=       CLIPTextModel(lowerCamelCase__				)

							__lowerCamelCase									=       XLMRobertaTokenizer.from_pretrained('hf-internal-testing/tiny-xlm-roberta'				)
							__lowerCamelCase									=       77

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


			def  lowercase_       ( self     ,		lowerCamelCase__     ,		lowerCamelCase__=0				)       -> Dict:



							'''simple docstring'''
							if str(lowerCamelCase__				).startswith('mps'				):
											__lowerCamelCase									=       torch.manual_seed(lowerCamelCase__				)
							else:
											__lowerCamelCase									=       torch.Generator(device=lowerCamelCase__				).manual_seed(lowerCamelCase__				)
							__lowerCamelCase									=       {
							    '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  lowercase_       ( self				)       -> Union[str, Any]:



							'''simple docstring'''
							super().test_attention_slicing_forward_pass(expected_max_diff=3e-3				)


			def  lowercase_       ( self				)       -> int:



							'''simple docstring'''
							super().test_inference_batch_single_identical(expected_max_diff=3e-3				)


			def  lowercase_       ( self				)       -> List[str]:



							'''simple docstring'''
							__lowerCamelCase									=       'cpu'  # ensure determinism for the device-dependent torch.Generator

							__lowerCamelCase									=       self.get_dummy_components()
							torch.manual_seed(0				)
							__lowerCamelCase									=       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=5_002     ,		)
							# TODO: remove after fixing the non-deterministic text encoder
							__lowerCamelCase									=       RobertaSeriesModelWithTransformation(lowerCamelCase__				)
							__lowerCamelCase									=       text_encoder

							__lowerCamelCase									=       AltDiffusionPipeline(**lowerCamelCase__				)
							__lowerCamelCase									=       alt_pipe.to(lowerCamelCase__				)
							alt_pipe.set_progress_bar_config(disable=lowerCamelCase__				)

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

							assert image.shape == (1, 64, 64, 3)
							__lowerCamelCase									=       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  lowercase_       ( self				)       -> Any:



							'''simple docstring'''
							__lowerCamelCase									=       'cpu'  # ensure determinism for the device-dependent torch.Generator

							__lowerCamelCase									=       self.get_dummy_components()
							__lowerCamelCase									=       PNDMScheduler(skip_prk_steps=lowerCamelCase__				)
							torch.manual_seed(0				)
							__lowerCamelCase									=       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=5_002     ,		)
							# TODO: remove after fixing the non-deterministic text encoder
							__lowerCamelCase									=       RobertaSeriesModelWithTransformation(lowerCamelCase__				)
							__lowerCamelCase									=       text_encoder
							__lowerCamelCase									=       AltDiffusionPipeline(**lowerCamelCase__				)
							__lowerCamelCase									=       alt_pipe.to(lowerCamelCase__				)
							alt_pipe.set_progress_bar_config(disable=lowerCamelCase__				)

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

							assert image.shape == (1, 64, 64, 3)
							__lowerCamelCase									=       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					):

			"""simple docstring"""







			def  lowercase_       ( self				)       -> str:



							'''simple docstring'''
							# clean up the VRAM after each test
							super().tearDown()
							gc.collect()
							torch.cuda.empty_cache()


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



							'''simple docstring'''
							# make sure here that pndm scheduler skips prk
							__lowerCamelCase									=       AltDiffusionPipeline.from_pretrained('BAAI/AltDiffusion'     ,		safety_checker=lowerCamelCase__				)
							__lowerCamelCase									=       alt_pipe.to(lowerCamelCase__				)
							alt_pipe.set_progress_bar_config(disable=lowerCamelCase__				)

							__lowerCamelCase									=       'A painting of a squirrel eating a burger'
							__lowerCamelCase									=       torch.manual_seed(0				)
							__lowerCamelCase									=       alt_pipe([prompt]     ,		generator=lowerCamelCase__     ,		guidance_scale=6.0     ,		num_inference_steps=20     ,		output_type='np'				)

							__lowerCamelCase									=       output.images

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

							assert image.shape == (1, 512, 512, 3)
							__lowerCamelCase									=       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  lowercase_       ( self				)       -> int:



							'''simple docstring'''
							__lowerCamelCase									=       DDIMScheduler.from_pretrained('BAAI/AltDiffusion'     ,		subfolder='scheduler'				)

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

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

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

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

							assert image.shape == (1, 512, 512, 3)
							__lowerCamelCase									=       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 __future__ import annotations

import math
from collections import Counter
from string import ascii_lowercase







def      lowerCamelCase_						(			UpperCamelCase__							: str )	->			None:





				"""simple docstring"""


				__lowerCamelCase			,		__lowerCamelCase									=       analyze_text(UpperCamelCase__ )
				__lowerCamelCase									=       list(' ' + ascii_lowercase )
				# what is our total sum of probabilities.
				__lowerCamelCase									=       sum(single_char_strings.values() )

				# one length string
				__lowerCamelCase									=       0
				# for each alpha we go in our dict and if it is in it we calculate entropy
				for ch in my_alphas:
								if ch in single_char_strings:
												__lowerCamelCase									=       single_char_strings[ch]
												__lowerCamelCase									=       my_str / all_sum
												my_fir_sum += prob * math.loga(UpperCamelCase__ )  # entropy formula.

    # print entropy
				print(F"""{round(-1 * my_fir_sum ):.1f}""" )

				# two len string
				__lowerCamelCase									=       sum(two_char_strings.values() )
				__lowerCamelCase									=       0
				# for each alpha (two in size) calculate entropy.
				for cha in my_alphas:
								for cha in my_alphas:
												__lowerCamelCase									=       cha + cha
												if sequence in two_char_strings:
																__lowerCamelCase									=       two_char_strings[sequence]
																__lowerCamelCase									=       int(UpperCamelCase__ ) / all_sum
																my_sec_sum += prob * math.loga(UpperCamelCase__ )

    # print second entropy
				print(F"""{round(-1 * my_sec_sum ):.1f}""" )

				# print the difference between them
				print(F"""{round((-1 * my_sec_sum) - (-1 * my_fir_sum) ):.1f}""" )







def      lowerCamelCase_						(			UpperCamelCase__							: str )	->			tuple[dict, dict]:





				"""simple docstring"""


				__lowerCamelCase									=       Counter()  # type: ignore
				__lowerCamelCase									=       Counter()  # type: ignore
				single_char_strings[text[-1]] += 1

				# first case when we have space at start.
				two_char_strings[" " + text[0]] += 1
				for i in range(0					,	len(UpperCamelCase__ ) - 1 ):
								single_char_strings[text[i]] += 1
								two_char_strings[text[i : i + 2]] += 1
				return single_char_strings, two_char_strings







def      lowerCamelCase_						(			)	->			Dict:





				"""simple docstring"""


				import doctest

				doctest.testmod()
				# text = (
				#     "Had repulsive dashwoods suspicion sincerity but advantage now him. Remark "
				#     "easily garret nor nay. Civil those mrs enjoy shy fat merry. You greatest "
				#     "jointure saw horrible. He private he on be imagine suppose. Fertile "
				#     "beloved evident through no service elderly is. Blind there if every no so "
				#     "at. Own neglected you preferred way sincerity delivered his attempted. To "
				#     "of message cottage windows do besides against uncivil.  Delightful "
				#     "unreserved impossible few estimating men favourable see entreaties. She "
				#     "propriety immediate was improving. He or entrance humoured likewise "
				#     "moderate. Much nor game son say feel. Fat make met can must form into "
				#     "gate. Me we offending prevailed discovery. "
				# )

				# calculate_prob(text)


if __name__ == "__main__":
					main()


