






'''simple docstring'''





import gc
import unittest

from diffusers import FlaxDPMSolverMultistepScheduler, FlaxStableDiffusionPipeline
from diffusers.utils import is_flax_available, slow
from diffusers.utils.testing_utils import require_flax


if is_flax_available():
					import jax
					import jax.numpy as jnp
					from flax.jax_utils import replicate
					from flax.training.common_utils import shard
@slow
@require_flax
class      UpperCamelCase (				unittest.TestCase ):







		"""simple docstring"""



		def       A(  self		:       Optional[int]		)    ->   Any:


						'''simple docstring'''
						super().tearDown()
						gc.collect()
		def       A(  self		:       Optional[Any]		)    ->   Optional[int]:


						'''simple docstring'''
						A ,		A	 =      FlaxStableDiffusionPipeline.from_pretrained(
						    'stabilityai/stable-diffusion-2'						,revision='bf16'						,dtype=jnp.bfloataa						,)

						A	 =      'A painting of a squirrel eating a burger'
						A	 =      jax.device_count()
						A	 =      num_samples * [prompt]
						A	 =      sd_pipe.prepare_inputs(_SCREAMING_SNAKE_CASE		)

						A	 =      replicate(_SCREAMING_SNAKE_CASE		)
						A	 =      shard(_SCREAMING_SNAKE_CASE		)

						A	 =      jax.random.PRNGKey(0		)
						A	 =      jax.random.split(_SCREAMING_SNAKE_CASE						,jax.device_count()		)

						A	 =      sd_pipe(_SCREAMING_SNAKE_CASE						,_SCREAMING_SNAKE_CASE						,_SCREAMING_SNAKE_CASE						,num_inference_steps=2_5						,jit=_SCREAMING_SNAKE_CASE		)[0]
						assert images.shape == (jax.device_count(), 1, 7_6_8, 7_6_8, 3)

						A	 =      images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]		)
						A	 =      images[0, 2_5_3:2_5_6, 2_5_3:2_5_6, -1]

						A	 =      jnp.asarray(jax.device_get(image_slice.flatten()		)		)
						A	 =      jnp.array([0.4238, 0.4414, 0.4395, 0.4453, 0.4629, 0.4590, 0.4531, 0.45508, 0.4512]		)
						print(f'output_slice: {output_slice}'		)
						assert jnp.abs(output_slice - expected_slice		).max() < 1E-2







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


						'''simple docstring'''
						A	 =      'stabilityai/stable-diffusion-2'
						A ,		A	 =      FlaxDPMSolverMultistepScheduler.from_pretrained(_SCREAMING_SNAKE_CASE						,subfolder='scheduler'		)
						A ,		A	 =      FlaxStableDiffusionPipeline.from_pretrained(
						    _SCREAMING_SNAKE_CASE						,scheduler=_SCREAMING_SNAKE_CASE						,revision='bf16'						,dtype=jnp.bfloataa						,)
						A	 =      scheduler_params

						A	 =      'A painting of a squirrel eating a burger'
						A	 =      jax.device_count()
						A	 =      num_samples * [prompt]
						A	 =      sd_pipe.prepare_inputs(_SCREAMING_SNAKE_CASE		)

						A	 =      replicate(_SCREAMING_SNAKE_CASE		)
						A	 =      shard(_SCREAMING_SNAKE_CASE		)

						A	 =      jax.random.PRNGKey(0		)
						A	 =      jax.random.split(_SCREAMING_SNAKE_CASE						,jax.device_count()		)

						A	 =      sd_pipe(_SCREAMING_SNAKE_CASE						,_SCREAMING_SNAKE_CASE						,_SCREAMING_SNAKE_CASE						,num_inference_steps=2_5						,jit=_SCREAMING_SNAKE_CASE		)[0]
						assert images.shape == (jax.device_count(), 1, 7_6_8, 7_6_8, 3)

						A	 =      images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]		)
						A	 =      images[0, 2_5_3:2_5_6, 2_5_3:2_5_6, -1]

						A	 =      jnp.asarray(jax.device_get(image_slice.flatten()		)		)
						A	 =      jnp.array([0.4336, 0.42969, 0.4453, 0.4199, 0.4297, 0.4531, 0.4434, 0.4434, 0.4297]		)
						print(f'output_slice: {output_slice}'		)
						assert jnp.abs(output_slice - expected_slice		).max() < 1E-2













def      snake_case		(							UpperCAmelCase	:       Optional[int],						UpperCAmelCase	:       Union[str, Any]   ):
				A	 =      ''
				for i in table:
								res += inp[i - 1]
				return res







def      snake_case		(							UpperCAmelCase	:       Union[str, Any]   ):
				return data[1:] + data[0]







def      snake_case		(							UpperCAmelCase	:       Union[str, Any],						UpperCAmelCase	:       Dict   ):
				A	 =      ''
				for i in range(len(UpperCAmelCase   )   ):
								if a[i] == b[i]:
												res += "0"
								else:
												res += "1"
				return res







def      snake_case		(							UpperCAmelCase	:       int,						UpperCAmelCase	:       Optional[Any]   ):
				A	 =      int('0b' + data[0] + data[-1],						2   )
				A	 =      int('0b' + data[1:3],						2   )
				return bin(s[row][col]   )[2:]







def      snake_case		(							UpperCAmelCase	:       Optional[Any],						UpperCAmelCase	:       Any,						UpperCAmelCase	:       Union[str, Any],						UpperCAmelCase	:       Optional[Any],						UpperCAmelCase	:       Optional[int]   ):
				A	 =      message[:4]
				A	 =      message[4:]
				A	 =      apply_table(UpperCAmelCase,						UpperCAmelCase   )
				A	 =      xor(UpperCAmelCase,						UpperCAmelCase   )
				A	 =      apply_sbox(UpperCAmelCase,						temp[:4]   )  # noqa: E741
				A	 =      apply_sbox(UpperCAmelCase,						temp[4:]   )
				A	 =      '0' * (2 - len(UpperCAmelCase   )) + l  # noqa: E741
				A	 =      '0' * (2 - len(UpperCAmelCase   )) + r
				A	 =      apply_table(l + r,						UpperCAmelCase   )
				A	 =      xor(UpperCAmelCase,						UpperCAmelCase   )
				return temp + right


if __name__ == "__main__":
					lowerCAmelCase_           = input('Enter 10 bit key: ')
					lowerCAmelCase_           = input('Enter 8 bit message: ')

					lowerCAmelCase_           = [6, 3, 7, 4, 8, 5, 10, 9]
					lowerCAmelCase_           = [3, 5, 2, 7, 4, 10, 1, 9, 8, 6]
					lowerCAmelCase_           = [2, 4, 3, 1]
					lowerCAmelCase_           = [2, 6, 3, 1, 4, 8, 5, 7]
					lowerCAmelCase_           = [4, 1, 3, 5, 7, 2, 8, 6]
					lowerCAmelCase_           = [4, 1, 2, 3, 2, 3, 4, 1]
					lowerCAmelCase_           = [[1, 0, 3, 2], [3, 2, 1, 0], [0, 2, 1, 3], [3, 1, 3, 2]]
					lowerCAmelCase_           = [[0, 1, 2, 3], [2, 0, 1, 3], [3, 0, 1, 0], [2, 1, 0, 3]]

					# key generation
					lowerCAmelCase_           = apply_table(key, paa_table)
					lowerCAmelCase_           = temp[:5]
					lowerCAmelCase_           = temp[5:]
					lowerCAmelCase_           = left_shift(left)
					lowerCAmelCase_           = left_shift(right)
					lowerCAmelCase_           = apply_table(left + right, pa_table)
					lowerCAmelCase_           = left_shift(left)
					lowerCAmelCase_           = left_shift(right)
					lowerCAmelCase_           = left_shift(left)
					lowerCAmelCase_           = left_shift(right)
					lowerCAmelCase_           = apply_table(left + right, pa_table)

					# encryption
					lowerCAmelCase_           = apply_table(message, IP)
					lowerCAmelCase_           = function(expansion, sa, sa, keya, temp)
					lowerCAmelCase_           = temp[4:] + temp[:4]
					lowerCAmelCase_           = function(expansion, sa, sa, keya, temp)
					lowerCAmelCase_           = apply_table(temp, IP_inv)
					print('Cipher text is:', CT)

					# decryption
					lowerCAmelCase_           = apply_table(CT, IP)
					lowerCAmelCase_           = function(expansion, sa, sa, keya, temp)
					lowerCAmelCase_           = temp[4:] + temp[:4]
					lowerCAmelCase_           = function(expansion, sa, sa, keya, temp)
					lowerCAmelCase_           = apply_table(temp, IP_inv)
					print('Plain text after decypting is:', PT)





