




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



				def       SCREAMING_SNAKE_CASE       (      self      :		Tuple      )  ->    Tuple:





								'''simple docstring'''




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



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





								'''simple docstring'''




								_lowerCAmelCase    :    Dict										=      FlaxStableDiffusionPipeline.from_pretrained(
								    """stabilityai/stable-diffusion-2"""		,	revision="""bf16"""		,	dtype=jnp.bfloataa		,	)

								_lowerCAmelCase    :    List[str]										=      """A painting of a squirrel eating a burger"""
								_lowerCAmelCase    :    List[str]										=      jax.device_count()
								_lowerCAmelCase    :    Dict										=      num_samples * [prompt]
								_lowerCAmelCase    :    List[Any]										=      sd_pipe.prepare_inputs(_UpperCAmelCase      )

								_lowerCAmelCase    :    int										=      replicate(_UpperCAmelCase      )
								_lowerCAmelCase    :    List[str]										=      shard(_UpperCAmelCase      )

								_lowerCAmelCase    :    str										=      jax.random.PRNGKey(0      )
								_lowerCAmelCase    :    List[Any]										=      jax.random.split(_UpperCAmelCase		,	jax.device_count()      )

								_lowerCAmelCase    :    List[Any]										=      sd_pipe(_UpperCAmelCase		,	_UpperCAmelCase		,	_UpperCAmelCase		,	num_inference_steps=25		,	jit=_UpperCAmelCase      )[0]
								assert images.shape == (jax.device_count(), 1, 768, 768, 3)

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

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





								'''simple docstring'''




								_lowerCAmelCase    :    List[str]										=      """stabilityai/stable-diffusion-2"""
								_lowerCAmelCase    :    List[Any]										=      FlaxDPMSolverMultistepScheduler.from_pretrained(_UpperCAmelCase		,	subfolder="""scheduler"""      )
								_lowerCAmelCase    :    Tuple										=      FlaxStableDiffusionPipeline.from_pretrained(
								    _UpperCAmelCase		,	scheduler=_UpperCAmelCase		,	revision="""bf16"""		,	dtype=jnp.bfloataa		,	)
								_lowerCAmelCase    :    Tuple										=      scheduler_params

								_lowerCAmelCase    :    Tuple										=      """A painting of a squirrel eating a burger"""
								_lowerCAmelCase    :    Optional[int]										=      jax.device_count()
								_lowerCAmelCase    :    int										=      num_samples * [prompt]
								_lowerCAmelCase    :    Any										=      sd_pipe.prepare_inputs(_UpperCAmelCase      )

								_lowerCAmelCase    :    List[str]										=      replicate(_UpperCAmelCase      )
								_lowerCAmelCase    :    Tuple										=      shard(_UpperCAmelCase      )

								_lowerCAmelCase    :    Union[str, Any]										=      jax.random.PRNGKey(0      )
								_lowerCAmelCase    :    int										=      jax.random.split(_UpperCAmelCase		,	jax.device_count()      )

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

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

								_lowerCAmelCase    :    Optional[Any]										=      jnp.asarray(jax.device_get(image_slice.flatten()      )      )
								_lowerCAmelCase    :    Union[str, Any]										=      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 							_UpperCAmelCase  (UpperCamelCase_ : str       ):
				'''simple docstring'''
				_lowerCAmelCase    :    Dict										=      [0] * len(UpperCamelCase_       )

				for i in range(1      ,    len(UpperCamelCase_       )       ):
								# use last results for better performance - dynamic programming
								_lowerCAmelCase    :    Union[str, Any]										=      prefix_result[i - 1]
								while j > 0 and input_string[i] != input_string[j]:
												_lowerCAmelCase    :    str										=      prefix_result[j - 1]

								if input_string[i] == input_string[j]:
												j += 1
								_lowerCAmelCase    :    Optional[int]										=      j

				return prefix_result




def 							_UpperCAmelCase  (UpperCamelCase_ : str       ):
				'''simple docstring'''
				return max(prefix_function(UpperCamelCase_       )       )


if __name__ == "__main__":
	import doctest

	doctest.testmod()

