



'''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 				__lowerCamelCase     (							self  :  Union[str, Any]):







											'''simple docstring'''

											super().tearDown()
											gc.collect()

				def 				__lowerCamelCase     (							self  :  Optional[Any]):







											'''simple docstring'''

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

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

											__lowercase       	=replicate(_lowerCAmelCase)
											__lowercase       	=shard(_lowerCAmelCase)

											__lowercase       	=jax.random.PRNGKey(0)
											__lowercase       	=jax.random.split(_lowerCAmelCase  ,       jax.device_count())

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

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

											__lowercase       	=jnp.asarray(jax.device_get(image_slice.flatten()))
											__lowercase       	=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 				__lowerCamelCase     (							self  :  List[Any]):







											'''simple docstring'''

											__lowercase       	='stabilityai/stable-diffusion-2'
											__lowercase	,				__lowercase       	=FlaxDPMSolverMultistepScheduler.from_pretrained(_lowerCAmelCase  ,       subfolder='scheduler')
											__lowercase	,				__lowercase       	=FlaxStableDiffusionPipeline.from_pretrained(
											    _lowerCAmelCase  ,       scheduler=_lowerCAmelCase  ,       revision='bf16'  ,       dtype=jnp.bfloataa  ,       )
											__lowercase       	=scheduler_params

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

											__lowercase       	=replicate(_lowerCAmelCase)
											__lowercase       	=shard(_lowerCAmelCase)

											__lowercase       	=jax.random.PRNGKey(0)
											__lowercase       	=jax.random.split(_lowerCAmelCase  ,       jax.device_count())

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

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

											__lowercase       	=jnp.asarray(jax.device_get(image_slice.flatten()))
											__lowercase       	=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'''



from functools import lru_cache



def   _A       (   _lowerCAmelCase					):



							"""simple docstring"""



							__lowercase       	=2
							__lowercase       	=set()
							while i * i <= n:
														if n % i:
																					i += 1
														else:
																					n //= i
																					factors.add(_lowerCAmelCase					)
							if n > 1:
														factors.add(_lowerCAmelCase					)
							return factors



@lru_cache
def   _A       (   _lowerCAmelCase					):



							"""simple docstring"""



							return len(unique_prime_factors(_lowerCAmelCase					)					)



def   _A       (   _lowerCAmelCase					):



							"""simple docstring"""



							return len(set(_lowerCAmelCase					)					) in (0, 1)



def   _A       (   _lowerCAmelCase					):



							"""simple docstring"""



							__lowercase       	=2

							while True:
														# Increment each value of a generated range
														__lowercase       	=[base + i for i in range(_lowerCAmelCase					)]

														# Run elements through out unique_prime_factors function
														# Append our target number to the end.
														__lowercase       	=[upf_len(_lowerCAmelCase					) for x in group]
														checker.append(_lowerCAmelCase					)

														# If all numbers in the list are equal, return the group variable.
														if equality(_lowerCAmelCase					):
																					return group

														# Increment our base variable by 1
														base += 1



def   _A       (   _lowerCAmelCase = 4					):



							"""simple docstring"""



							__lowercase       	=run(_lowerCAmelCase					)
							return results[0] if len(_lowerCAmelCase					) else None


if __name__ == "__main__":
		print(solution())



