




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


						def   __lowercase (							self				:			Any					):
										# clean up the VRAM after each test
										super().tearDown()
										gc.collect()


						def   __lowercase (							self				:			List[Any]					):
										lowerCAmelCase ,						lowerCAmelCase  =				FlaxStableDiffusionPipeline.from_pretrained(
										    """stabilityai/stable-diffusion-2"""						,						revision="""bf16"""						,						dtype=jnp.bfloataa						,						)

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

										lowerCAmelCase  =				replicate(lowerCAmelCase					)
										lowerCAmelCase  =				shard(lowerCAmelCase					)

										lowerCAmelCase  =				jax.random.PRNGKey(0					)
										lowerCAmelCase  =				jax.random.split(lowerCAmelCase						,						jax.device_count()					)

										lowerCAmelCase  =				sd_pipe(lowerCAmelCase						,						lowerCAmelCase						,						lowerCAmelCase						,						num_inference_steps=25						,						jit=lowerCAmelCase					)[0]
										assert images.shape == (jax.device_count(), 1, 768, 768, 3)

										lowerCAmelCase  =				images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]					)
										lowerCAmelCase  =				images[0, 253:256, 253:256, -1]

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




						def   __lowercase (							self				:			Tuple					):
										lowerCAmelCase  =				"""stabilityai/stable-diffusion-2"""
										lowerCAmelCase ,						lowerCAmelCase  =				FlaxDPMSolverMultistepScheduler.from_pretrained(lowerCAmelCase						,						subfolder="""scheduler"""					)
										lowerCAmelCase ,						lowerCAmelCase  =				FlaxStableDiffusionPipeline.from_pretrained(
										    lowerCAmelCase						,						scheduler=lowerCAmelCase						,						revision="""bf16"""						,						dtype=jnp.bfloataa						,						)
										lowerCAmelCase  =				scheduler_params

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

										lowerCAmelCase  =				replicate(lowerCAmelCase					)
										lowerCAmelCase  =				shard(lowerCAmelCase					)

										lowerCAmelCase  =				jax.random.PRNGKey(0					)
										lowerCAmelCase  =				jax.random.split(lowerCAmelCase						,						jax.device_count()					)

										lowerCAmelCase  =				sd_pipe(lowerCAmelCase						,						lowerCAmelCase						,						lowerCAmelCase						,						num_inference_steps=25						,						jit=lowerCAmelCase					)[0]
										assert images.shape == (jax.device_count(), 1, 768, 768, 3)

										lowerCAmelCase  =				images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]					)
										lowerCAmelCase  =				images[0, 253:256, 253:256, -1]

										lowerCAmelCase  =				jnp.asarray(jax.device_get(image_slice.flatten()					)					)
										lowerCAmelCase  =				jnp.array([0.4336, 0.4_2969, 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











"""simple docstring"""







import os
import string
import sys


a      =							1 << 8

a      =							{
    'tab': ord('\t'),
    'newline': ord('\r'),
    'esc': 2_7,
    'up': 6_5 + ARROW_KEY_FLAG,
    'down': 6_6 + ARROW_KEY_FLAG,
    'right': 6_7 + ARROW_KEY_FLAG,
    'left': 6_8 + ARROW_KEY_FLAG,
    'mod_int': 9_1,
    'undefined': sys.maxsize,
    'interrupt': 3,
    'insert': 5_0,
    'delete': 5_1,
    'pg_up': 5_3,
    'pg_down': 5_4,
}

a      =							KEYMAP['up']
a      =							KEYMAP['left']

if sys.platform == "win32":
		a      =							[]
		a      =							{
		    b'\xe0H': KEYMAP['up'] - ARROW_KEY_FLAG,
		    b'\x00H': KEYMAP['up'] - ARROW_KEY_FLAG,
		    b'\xe0P': KEYMAP['down'] - ARROW_KEY_FLAG,
		    b'\x00P': KEYMAP['down'] - ARROW_KEY_FLAG,
		    b'\xe0M': KEYMAP['right'] - ARROW_KEY_FLAG,
		    b'\x00M': KEYMAP['right'] - ARROW_KEY_FLAG,
		    b'\xe0K': KEYMAP['left'] - ARROW_KEY_FLAG,
		    b'\x00K': KEYMAP['left'] - ARROW_KEY_FLAG,
		}

for i in range(1_0):
		a      =							ord(str(i))





def 							lowercase					()     ->  str:


				'''simple docstring'''






				if os.name == "nt":
								import msvcrt

								lowerCAmelCase  =				"""mbcs"""
								# Flush the keyboard buffer
								while msvcrt.kbhit():
												msvcrt.getch()
								if len(snake_case__ ) == 0:
												# Read the keystroke
												lowerCAmelCase  =				msvcrt.getch()

												# If it is a prefix char, get second part
												if ch in (b"\x00", b"\xe0"):
																lowerCAmelCase  =				ch + msvcrt.getch()
																# Translate actual Win chars to bullet char types
																try:
																				lowerCAmelCase  =				chr(WIN_KEYMAP[cha] )
																				WIN_CH_BUFFER.append(chr(KEYMAP["""mod_int"""] ) )
																				WIN_CH_BUFFER.append(snake_case__ )
																				if ord(snake_case__ ) in (
																				    KEYMAP["insert"] - 1 << 9,
																				    KEYMAP["delete"] - 1 << 9,
																				    KEYMAP["pg_up"] - 1 << 9,
																				    KEYMAP["pg_down"] - 1 << 9,
																				):
																								WIN_CH_BUFFER.append(chr(126 ) )
																				lowerCAmelCase  =				chr(KEYMAP["""esc"""] )
																except KeyError:
																				lowerCAmelCase  =				cha[1]
												else:
																lowerCAmelCase  =				ch.decode(snake_case__ )
								else:
												lowerCAmelCase  =				WIN_CH_BUFFER.pop(0 )
				elif os.name == "posix":
								import termios
								import tty

								lowerCAmelCase  =				sys.stdin.fileno()
								lowerCAmelCase  =				termios.tcgetattr(snake_case__ )
								try:
												tty.setraw(snake_case__ )
												lowerCAmelCase  =				sys.stdin.read(1 )
								finally:
												termios.tcsetattr(snake_case__      ,       termios.TCSADRAIN      ,       snake_case__ )
				return ch






def 							lowercase					()     ->  List[str]:


				'''simple docstring'''






				lowerCAmelCase  =				get_raw_chars()
				if ord(snake_case__ ) in [KEYMAP["interrupt"], KEYMAP["newline"]]:
								return char

				elif ord(snake_case__ ) == KEYMAP["esc"]:
								lowerCAmelCase  =				get_raw_chars()
								if ord(snake_case__ ) == KEYMAP["mod_int"]:
												lowerCAmelCase  =				get_raw_chars()
												if ord(snake_case__ ) >= KEYMAP["arrow_begin"] - ARROW_KEY_FLAG and ord(snake_case__ ) <= KEYMAP["arrow_end"] - ARROW_KEY_FLAG:
																return chr(ord(snake_case__ ) + ARROW_KEY_FLAG )
												else:
																return KEYMAP["undefined"]
								else:
												return get_raw_chars()

				else:
								if char in string.printable:
												return char
								else:
												return KEYMAP["undefined"]





