import gc
import unittest

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

from diffusers import (
    AutoencoderKL,
    DiffusionPipeline,
    EulerDiscreteScheduler,
    StableDiffusionXLPipeline,
    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 _lowercase					(       _UpperCAmelCase   ,				_UpperCAmelCase   ,				unittest.TestCase						):





       """simple docstring"""




       lowerCAmelCase__       =					StableDiffusionXLPipeline
       lowerCAmelCase__       =					TEXT_TO_IMAGE_PARAMS
       lowerCAmelCase__       =					TEXT_TO_IMAGE_BATCH_PARAMS
       lowerCAmelCase__       =					TEXT_TO_IMAGE_IMAGE_PARAMS
       lowerCAmelCase__       =					TEXT_TO_IMAGE_IMAGE_PARAMS



       def     _UpperCAmelCase       (    self  ):







              '''simple docstring'''






              torch.manual_seed(0  )
              _lowercase				=	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""")     ,   attention_head_dim=(2, 4)     ,   use_linear_projection=UpperCAmelCase     ,   addition_embed_type="""text_time"""     ,   addition_time_embed_dim=8     ,   transformer_layers_per_block=(1, 2)     ,   projection_class_embeddings_input_dim=80     ,   cross_attention_dim=64     ,   )
              _lowercase				=	EulerDiscreteScheduler(
                  beta_start=0.00_085     ,   beta_end=0.012     ,   steps_offset=1     ,   beta_schedule="""scaled_linear"""     ,   timestep_spacing="""leading"""     ,   )
              torch.manual_seed(0  )
              _lowercase				=	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     ,   sample_size=128     ,   )
              torch.manual_seed(0  )
              _lowercase				=	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=1000     ,   hidden_act="""gelu"""     ,   projection_dim=32     ,   )
              _lowercase				=	CLIPTextModel(UpperCAmelCase  )
              _lowercase				=	CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""     ,   local_files_only=UpperCAmelCase  )

              _lowercase				=	CLIPTextModelWithProjection(UpperCAmelCase  )
              _lowercase				=	CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""     ,   local_files_only=UpperCAmelCase  )

              _lowercase				=	{
                  """unet""": unet,
                  """scheduler""": scheduler,
                  """vae""": vae,
                  """text_encoder""": text_encoder,
                  """tokenizer""": tokenizer,
                  """text_encoder_2""": text_encoder_a,
                  """tokenizer_2""": tokenizer_a,
                  # "safety_checker": None,
                  # "feature_extractor": None,
              }
              return components



       def     _UpperCAmelCase       (    self     ,   UpperCAmelCase     ,   UpperCAmelCase=0  ):







              '''simple docstring'''






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



       def     _UpperCAmelCase       (    self  ):







              '''simple docstring'''






              _lowercase				=	"""cpu"""  # ensure determinism for the device-dependent torch.Generator
              _lowercase				=	self.get_dummy_components()
              _lowercase				=	StableDiffusionXLPipeline(**UpperCAmelCase  )
              _lowercase				=	sd_pipe.to(UpperCAmelCase  )
              sd_pipe.set_progress_bar_config(disable=UpperCAmelCase  )

              _lowercase				=	self.get_dummy_inputs(UpperCAmelCase  )
              _lowercase				=	sd_pipe(**UpperCAmelCase  ).images
              _lowercase				=	image[0, -3:, -3:, -1]

              assert image.shape == (1, 64, 64, 3)
              _lowercase				=	np.array([0.5_873, 0.6_128, 0.4_797, 0.5_122, 0.5_674, 0.4_639, 0.5_227, 0.5_149, 0.4_747]  )

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



       def     _UpperCAmelCase       (    self  ):







              '''simple docstring'''






              _lowercase				=	self.get_dummy_components()
              _lowercase				=	StableDiffusionXLPipeline(**UpperCAmelCase  )
              _lowercase				=	sd_pipe.to(UpperCAmelCase  )
              _lowercase				=	sd_pipe.to(UpperCAmelCase  )
              sd_pipe.set_progress_bar_config(disable=UpperCAmelCase  )

              # forward without prompt embeds
              _lowercase				=	self.get_dummy_inputs(UpperCAmelCase  )
              _lowercase				=	3 * ["""this is a negative prompt"""]
              _lowercase				=	negative_prompt
              _lowercase				=	3 * [inputs["""prompt"""]]

              _lowercase				=	sd_pipe(**UpperCAmelCase  )
              _lowercase				=	output.images[0, -3:, -3:, -1]

              # forward with prompt embeds
              _lowercase				=	self.get_dummy_inputs(UpperCAmelCase  )
              _lowercase				=	3 * ["""this is a negative prompt"""]
              _lowercase				=	3 * [inputs.pop("""prompt"""  )]

              (
                  (
                  _lowercase       
              )				,       (
                  _lowercase       
              )				,       (
                  _lowercase       
              )				,       (
                  _lowercase       
              )				,              
              )				=	sd_pipe.encode_prompt(UpperCAmelCase     ,   negative_prompt=UpperCAmelCase  )

              _lowercase				=	sd_pipe(
                  **UpperCAmelCase     ,   prompt_embeds=UpperCAmelCase     ,   negative_prompt_embeds=UpperCAmelCase     ,   pooled_prompt_embeds=UpperCAmelCase     ,   negative_pooled_prompt_embeds=UpperCAmelCase     ,   )
              _lowercase				=	output.images[0, -3:, -3:, -1]

              # make sure that it's equal
              assert np.abs(image_slice_a.flatten() - image_slice_a.flatten()  ).max() < 1e-4



       def     _UpperCAmelCase       (    self  ):







              '''simple docstring'''






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



       def     _UpperCAmelCase       (    self  ):







              '''simple docstring'''






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






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





       """simple docstring"""




       def     _UpperCAmelCase       (    self  ):







              '''simple docstring'''






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



       def     _UpperCAmelCase       (    self     ,   UpperCAmelCase     ,   UpperCAmelCase="cpu"     ,   UpperCAmelCase=torch.floataa     ,   UpperCAmelCase=0  ):







              '''simple docstring'''






              _lowercase				=	torch.Generator(device=UpperCAmelCase  ).manual_seed(UpperCAmelCase  )
              _lowercase				=	np.random.RandomState(UpperCAmelCase  ).standard_normal((1, 4, 64, 64)  )
              _lowercase				=	torch.from_numpy(UpperCAmelCase  ).to(device=UpperCAmelCase     ,   dtype=UpperCAmelCase  )
              _lowercase				=	{
                  """prompt""": """a photograph of an astronaut riding a horse""",
                  """latents""": latents,
                  """generator""": generator,
                  """num_inference_steps""": 3,
                  """guidance_scale""": 7.5,
                  """output_type""": """numpy""",
              }
              return inputs



       def     _UpperCAmelCase       (    self  ):







              '''simple docstring'''






              _lowercase				=	DiffusionPipeline.from_pretrained("""stabilityai/stable-diffusion-2-base"""  )
              pipe.to(UpperCAmelCase  )
              pipe.set_progress_bar_config(disable=UpperCAmelCase  )

              _lowercase				=	self.get_inputs(UpperCAmelCase  )
              _lowercase				=	pipe(**UpperCAmelCase  ).images
              _lowercase				=	image[0, -3:, -3:, -1].flatten()

              assert image.shape == (1, 512, 512, 3)
              _lowercase				=	np.array([0.49_493, 0.47_896, 0.40_798, 0.54_214, 0.53_212, 0.48_202, 0.47_656, 0.46_329, 0.48_506]  )
              assert np.abs(image_slice - expected_slice  ).max() < 7e-3



import numpy as np


def 					__lowerCAmelCase			(  _A ,_A="utf-8" ,_A="surrogatepass"					):



							"""simple docstring"""


							_lowercase				=	bin(int.from_bytes(text.encode(_A ,_A					) ,"""big"""					)					)[2:]
							return bits.zfill(8 * ((len(_A					) + 7) // 8)					)


def 					__lowerCAmelCase			(  _A ,_A="utf-8" ,_A="surrogatepass"					):



							"""simple docstring"""


							_lowercase				=	int(_A ,2					)
							return n.to_bytes((n.bit_length() + 7) // 8 ,"""big"""					).decode(_A ,_A					) or "\0"


def 					__lowerCAmelCase			(  _A ,_A					):



							"""simple docstring"""


							if size_par + len(_A					) <= 2**size_par - (len(_A					) - 1):
														raise ValueError("""size of parity don't match with size of data"""					)

							_lowercase				=	[]
							_lowercase				=	[]
							_lowercase				=	[bin(_A					)[2:] for x in range(1 ,size_par + len(_A					) + 1					)]

							# sorted information data for the size of the output data
							_lowercase				=	[]
							# data position template + parity
							_lowercase				=	[]
							# parity bit counter
							_lowercase				=	0
							# counter position of data bits
							_lowercase				=	0

							for x in range(1 ,size_par + len(_A					) + 1					):
														# Performs a template of bit positions - who should be given,
														# and who should be parity
														if qtd_bp < size_par:
																					if (np.log(_A					) / np.log(2					)).is_integer():
																												data_out_gab.append("""P"""					)
																												_lowercase				=	qtd_bp + 1
																					else:
																												data_out_gab.append("""D"""					)
														else:
																					data_out_gab.append("""D"""					)

														# Sorts the data to the new output size
														if data_out_gab[-1] == "D":
																					data_ord.append(data[cont_data]					)
																					cont_data += 1
														else:
																					data_ord.append(_A					)

    # Calculates parity
							_lowercase				=	0  # parity bit counter
							for bp in range(1 ,size_par + 1					):
														# Bit counter one for a given parity
														_lowercase				=	0
														# counter to control the loop reading
														_lowercase				=	0
														for x in data_ord:
																					if x is not None:
																												try:
																																			_lowercase				=	(bin_pos[cont_loop])[-1 * (bp)]
																												except IndexError:
																																			_lowercase				=	"""0"""
																												if aux == "1" and x == "1":
																																			cont_bo += 1
																					cont_loop += 1
														parity.append(cont_bo % 2					)

														qtd_bp += 1

							# Mount the message
							_lowercase				=	0  # parity bit counter
							for x in range(0 ,size_par + len(_A					)					):
														if data_ord[x] is None:
																					data_out.append(str(parity[cont_bp]					)					)
																					cont_bp += 1
														else:
																					data_out.append(data_ord[x]					)

							return data_out


def 					__lowerCAmelCase			(  _A ,_A					):



							"""simple docstring"""


							_lowercase				=	[]
							# Parity bit counter
							_lowercase				=	0
							# Counter p data bit reading
							_lowercase				=	0
							# list of parity received
							_lowercase				=	[]
							_lowercase				=	[]

							for x in range(1 ,len(_A					) + 1					):
														# Performs a template of bit positions - who should be given,
														#  and who should be parity
														if qtd_bp < size_par and (np.log(_A					) / np.log(2					)).is_integer():
																					data_out_gab.append("""P"""					)
																					_lowercase				=	qtd_bp + 1
														else:
																					data_out_gab.append("""D"""					)

														# Sorts the data to the new output size
														if data_out_gab[-1] == "D":
																					data_output.append(data[cont_data]					)
														else:
																					parity_received.append(data[cont_data]					)
														cont_data += 1

							# -----------calculates the parity with the data
							_lowercase				=	[]
							_lowercase				=	[]
							_lowercase				=	[bin(_A					)[2:] for x in range(1 ,size_par + len(_A					) + 1					)]

							#  sorted information data for the size of the output data
							_lowercase				=	[]
							# Data position feedback + parity
							_lowercase				=	[]
							# Parity bit counter
							_lowercase				=	0
							# Counter p data bit reading
							_lowercase				=	0

							for x in range(1 ,size_par + len(_A					) + 1					):
														# Performs a template position of bits - who should be given,
														# and who should be parity
														if qtd_bp < size_par and (np.log(_A					) / np.log(2					)).is_integer():
																					data_out_gab.append("""P"""					)
																					_lowercase				=	qtd_bp + 1
														else:
																					data_out_gab.append("""D"""					)

														# Sorts the data to the new output size
														if data_out_gab[-1] == "D":
																					data_ord.append(data_output[cont_data]					)
																					cont_data += 1
														else:
																					data_ord.append(_A					)

    # Calculates parity
							_lowercase				=	0  # parity bit counter
							for bp in range(1 ,size_par + 1					):
														# Bit counter one for a certain parity
														_lowercase				=	0
														# Counter to control loop reading
														_lowercase				=	0
														for x in data_ord:
																					if x is not None:
																												try:
																																			_lowercase				=	(bin_pos[cont_loop])[-1 * (bp)]
																												except IndexError:
																																			_lowercase				=	"""0"""
																												if aux == "1" and x == "1":
																																			cont_bo += 1
																					cont_loop += 1
														parity.append(str(cont_bo % 2					)					)

														qtd_bp += 1

							# Mount the message
							_lowercase				=	0  # Parity bit counter
							for x in range(0 ,size_par + len(_A					)					):
														if data_ord[x] is None:
																					data_out.append(str(parity[cont_bp]					)					)
																					cont_bp += 1
														else:
																					data_out.append(data_ord[x]					)

							_lowercase				=	parity_received == parity
							return data_output, ack


