



"""simple docstring"""





import os
import tempfile
import unittest

import numpy as np

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


if is_flax_available():
					import jax
					import jax.numpy as jnp
					from flax.jax_utils import replicate
					from flax.training.common_utils import shard

					from diffusers import FlaxDDIMScheduler, FlaxDiffusionPipeline, FlaxStableDiffusionPipeline



@require_flax
class  _lowerCAmelCase     (					unittest.TestCase		):





		"""simple docstring"""






		def  UpperCAmelCase__				(       self )			->    List[str]:







								'''simple docstring'''

								with tempfile.TemporaryDirectory() as tmpdirname:
														# pipeline has Flax weights
														snake_case_ :			Optional[Any]            =      FlaxDiffusionPipeline.from_pretrained(
														    """hf-internal-testing/tiny-stable-diffusion-pipe"""				,     safety_checker=__UpperCAmelCase				,     cache_dir=__UpperCAmelCase )

														snake_case_ :			Optional[Any]            =      [t[-1] for t in os.walk(os.path.join(__UpperCAmelCase				,     os.listdir(__UpperCAmelCase )[0]				,     """snapshots""" ) )]
														snake_case_ :			Dict            =      [item for sublist in all_root_files for item in sublist]

														# None of the downloaded files should be a PyTorch file even if we have some here:
														# https://huggingface.co/hf-internal-testing/tiny-stable-diffusion-pipe/blob/main/unet/diffusion_pytorch_model.bin
														assert not any(f.endswith(""".bin""" ) for f in files )




@slow
@require_flax
class  _lowerCAmelCase     (					unittest.TestCase		):





		"""simple docstring"""






		def  UpperCAmelCase__				(       self )			->    Tuple:







								'''simple docstring'''

								snake_case_ :			int            =      FlaxStableDiffusionPipeline.from_pretrained(
								    """hf-internal-testing/tiny-stable-diffusion-pipe"""				,     safety_checker=__UpperCAmelCase )

								snake_case_ :			Optional[Any]            =      (
								    "A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of"
								    " field, close up, split lighting, cinematic"
								)

								snake_case_ :			Tuple            =      jax.random.PRNGKey(0 )
								snake_case_ :			Optional[int]            =      4

								snake_case_ :			Optional[Any]            =      jax.device_count()
								snake_case_ :			Tuple            =      num_samples * [prompt]
								snake_case_ :			Dict            =      pipeline.prepare_inputs(__UpperCAmelCase )

								# shard inputs and rng
								snake_case_ :			List[str]            =      replicate(__UpperCAmelCase )
								snake_case_ :			Dict            =      jax.random.split(__UpperCAmelCase				,     __UpperCAmelCase )
								snake_case_ :			Tuple            =      shard(__UpperCAmelCase )

								snake_case_ :			str            =      pipeline(__UpperCAmelCase				,     __UpperCAmelCase				,     __UpperCAmelCase				,     __UpperCAmelCase				,     jit=__UpperCAmelCase ).images

								assert images.shape == (num_samples, 1, 6_4, 6_4, 3)
								if jax.device_count() == 8:
														assert np.abs(np.abs(images[0, 0, :2, :2, -2:]				,     dtype=np.floataa ).sum() - 4.151_4745 ) < 1E-3
														assert np.abs(np.abs(__UpperCAmelCase				,     dtype=np.floataa ).sum() - 4_9_9_4_7.8_7_5 ) < 5E-1

								snake_case_ :			int            =      pipeline.numpy_to_pil(np.asarray(images.reshape((num_samples,) + images.shape[-3:] ) ) )
								assert len(__UpperCAmelCase ) == num_samples





		def  UpperCAmelCase__				(       self )			->    List[Any]:







								'''simple docstring'''

								snake_case_ :			List[Any]            =      FlaxStableDiffusionPipeline.from_pretrained(
								    """CompVis/stable-diffusion-v1-4"""				,     revision="""flax"""				,     safety_checker=__UpperCAmelCase )

								snake_case_ :			Optional[int]            =      (
								    "A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of"
								    " field, close up, split lighting, cinematic"
								)

								snake_case_ :			Optional[int]            =      jax.random.PRNGKey(0 )
								snake_case_ :			Optional[int]            =      5_0

								snake_case_ :			Optional[int]            =      jax.device_count()
								snake_case_ :			Optional[int]            =      num_samples * [prompt]
								snake_case_ :			str            =      pipeline.prepare_inputs(__UpperCAmelCase )

								# shard inputs and rng
								snake_case_ :			Optional[Any]            =      replicate(__UpperCAmelCase )
								snake_case_ :			int            =      jax.random.split(__UpperCAmelCase				,     __UpperCAmelCase )
								snake_case_ :			List[str]            =      shard(__UpperCAmelCase )

								snake_case_ :			List[str]            =      pipeline(__UpperCAmelCase				,     __UpperCAmelCase				,     __UpperCAmelCase				,     __UpperCAmelCase				,     jit=__UpperCAmelCase ).images

								assert images.shape == (num_samples, 1, 5_1_2, 5_1_2, 3)
								if jax.device_count() == 8:
														assert np.abs((np.abs(images[0, 0, :2, :2, -2:]				,     dtype=np.floataa ).sum() - 0.0565_2401) ) < 1E-3
														assert np.abs((np.abs(__UpperCAmelCase				,     dtype=np.floataa ).sum() - 2_3_8_3_8_0_8.2) ) < 5E-1





		def  UpperCAmelCase__				(       self )			->    Optional[Any]:







								'''simple docstring'''

								snake_case_ :			Optional[int]            =      FlaxStableDiffusionPipeline.from_pretrained(
								    """CompVis/stable-diffusion-v1-4"""				,     revision="""bf16"""				,     dtype=jnp.bfloataa				,     safety_checker=__UpperCAmelCase )

								snake_case_ :			Tuple            =      (
								    "A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of"
								    " field, close up, split lighting, cinematic"
								)

								snake_case_ :			Optional[Any]            =      jax.random.PRNGKey(0 )
								snake_case_ :			List[Any]            =      5_0

								snake_case_ :			Union[str, Any]            =      jax.device_count()
								snake_case_ :			int            =      num_samples * [prompt]
								snake_case_ :			Optional[int]            =      pipeline.prepare_inputs(__UpperCAmelCase )

								# shard inputs and rng
								snake_case_ :			List[str]            =      replicate(__UpperCAmelCase )
								snake_case_ :			str            =      jax.random.split(__UpperCAmelCase				,     __UpperCAmelCase )
								snake_case_ :			Dict            =      shard(__UpperCAmelCase )

								snake_case_ :			Any            =      pipeline(__UpperCAmelCase				,     __UpperCAmelCase				,     __UpperCAmelCase				,     __UpperCAmelCase				,     jit=__UpperCAmelCase ).images

								assert images.shape == (num_samples, 1, 5_1_2, 5_1_2, 3)
								if jax.device_count() == 8:
														assert np.abs((np.abs(images[0, 0, :2, :2, -2:]				,     dtype=np.floataa ).sum() - 0.0400_3906) ) < 1E-3
														assert np.abs((np.abs(__UpperCAmelCase				,     dtype=np.floataa ).sum() - 2_3_7_3_5_1_6.7_5) ) < 5E-1





		def  UpperCAmelCase__				(       self )			->    Union[str, Any]:







								'''simple docstring'''

								snake_case_ :			Optional[int]            =      FlaxStableDiffusionPipeline.from_pretrained(
								    """CompVis/stable-diffusion-v1-4"""				,     revision="""bf16"""				,     dtype=jnp.bfloataa )

								snake_case_ :			Optional[Any]            =      (
								    "A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of"
								    " field, close up, split lighting, cinematic"
								)

								snake_case_ :			List[str]            =      jax.random.PRNGKey(0 )
								snake_case_ :			Dict            =      5_0

								snake_case_ :			Any            =      jax.device_count()
								snake_case_ :			Dict            =      num_samples * [prompt]
								snake_case_ :			int            =      pipeline.prepare_inputs(__UpperCAmelCase )

								# shard inputs and rng
								snake_case_ :			Optional[int]            =      replicate(__UpperCAmelCase )
								snake_case_ :			Union[str, Any]            =      jax.random.split(__UpperCAmelCase				,     __UpperCAmelCase )
								snake_case_ :			Optional[int]            =      shard(__UpperCAmelCase )

								snake_case_ :			Optional[Any]            =      pipeline(__UpperCAmelCase				,     __UpperCAmelCase				,     __UpperCAmelCase				,     __UpperCAmelCase				,     jit=__UpperCAmelCase ).images

								assert images.shape == (num_samples, 1, 5_1_2, 5_1_2, 3)
								if jax.device_count() == 8:
														assert np.abs((np.abs(images[0, 0, :2, :2, -2:]				,     dtype=np.floataa ).sum() - 0.0400_3906) ) < 1E-3
														assert np.abs((np.abs(__UpperCAmelCase				,     dtype=np.floataa ).sum() - 2_3_7_3_5_1_6.7_5) ) < 5E-1





		def  UpperCAmelCase__				(       self )			->    Any:







								'''simple docstring'''

								snake_case_ :			int            =      FlaxDDIMScheduler(
								    beta_start=0.0_0085				,     beta_end=0.012				,     beta_schedule="""scaled_linear"""				,     set_alpha_to_one=__UpperCAmelCase				,     steps_offset=1				,     )

								snake_case_ :			Optional[int]            =      FlaxStableDiffusionPipeline.from_pretrained(
								    """CompVis/stable-diffusion-v1-4"""				,     revision="""bf16"""				,     dtype=jnp.bfloataa				,     scheduler=__UpperCAmelCase				,     safety_checker=__UpperCAmelCase				,     )
								snake_case_ :			Union[str, Any]            =      scheduler.create_state()

								snake_case_ :			Tuple            =      scheduler_state

								snake_case_ :			Any            =      (
								    "A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of"
								    " field, close up, split lighting, cinematic"
								)

								snake_case_ :			Any            =      jax.random.PRNGKey(0 )
								snake_case_ :			List[Any]            =      5_0

								snake_case_ :			str            =      jax.device_count()
								snake_case_ :			Any            =      num_samples * [prompt]
								snake_case_ :			Optional[int]            =      pipeline.prepare_inputs(__UpperCAmelCase )

								# shard inputs and rng
								snake_case_ :			Dict            =      replicate(__UpperCAmelCase )
								snake_case_ :			Optional[Any]            =      jax.random.split(__UpperCAmelCase				,     __UpperCAmelCase )
								snake_case_ :			List[str]            =      shard(__UpperCAmelCase )

								snake_case_ :			Optional[Any]            =      pipeline(__UpperCAmelCase				,     __UpperCAmelCase				,     __UpperCAmelCase				,     __UpperCAmelCase				,     jit=__UpperCAmelCase ).images

								assert images.shape == (num_samples, 1, 5_1_2, 5_1_2, 3)
								if jax.device_count() == 8:
														assert np.abs((np.abs(images[0, 0, :2, :2, -2:]				,     dtype=np.floataa ).sum() - 0.0_4504_3945) ) < 1E-3
														assert np.abs((np.abs(__UpperCAmelCase				,     dtype=np.floataa ).sum() - 2_3_4_7_6_9_3.5) ) < 5E-1





		def  UpperCAmelCase__				(       self )			->    str:







								'''simple docstring'''

								snake_case_ :			Dict            =      (
								    "A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of"
								    " field, close up, split lighting, cinematic"
								)

								snake_case_ :			str            =      jax.device_count()
								snake_case_ :			Tuple            =      num_samples * [prompt]
								snake_case_ :			Union[str, Any]            =      jax.random.split(jax.random.PRNGKey(0 )				,     __UpperCAmelCase )

								snake_case_ :			Any            =      FlaxStableDiffusionPipeline.from_pretrained(
								    """CompVis/stable-diffusion-v1-4"""				,     revision="""bf16"""				,     dtype=jnp.bfloataa				,     safety_checker=__UpperCAmelCase				,     )

								snake_case_ :			Any            =      replicate(__UpperCAmelCase )
								snake_case_ :			Optional[Any]            =      pipeline.prepare_inputs(__UpperCAmelCase )
								snake_case_ :			Union[str, Any]            =      shard(__UpperCAmelCase )
								snake_case_ :			Optional[Any]            =      pipeline(__UpperCAmelCase				,     __UpperCAmelCase				,     __UpperCAmelCase				,     jit=__UpperCAmelCase ).images
								assert images.shape == (num_samples, 1, 5_1_2, 5_1_2, 3)
								snake_case_ :			Dict            =      images[2, 0, 2_5_6, 1_0:1_7, 1]

								# With memory efficient attention
								snake_case_ :			Optional[Any]            =      FlaxStableDiffusionPipeline.from_pretrained(
								    """CompVis/stable-diffusion-v1-4"""				,     revision="""bf16"""				,     dtype=jnp.bfloataa				,     safety_checker=__UpperCAmelCase				,     use_memory_efficient_attention=__UpperCAmelCase				,     )

								snake_case_ :			Optional[Any]            =      replicate(__UpperCAmelCase )
								snake_case_ :			Optional[int]            =      pipeline.prepare_inputs(__UpperCAmelCase )
								snake_case_ :			Dict            =      shard(__UpperCAmelCase )
								snake_case_ :			str            =      pipeline(__UpperCAmelCase				,     __UpperCAmelCase				,     __UpperCAmelCase				,     jit=__UpperCAmelCase ).images
								assert images_eff.shape == (num_samples, 1, 5_1_2, 5_1_2, 3)
								snake_case_ :			int            =      images[2, 0, 2_5_6, 1_0:1_7, 1]

								# I checked the results visually and they are very similar. However, I saw that the max diff is `1` and the `sum`
								# over the 8 images is exactly `256`, which is very suspicious. Testing a random slice for now.
								assert abs(slice_eff - slice ).max() < 1E-2









def     a_   (__A      )				->						Dict:



					"""simple docstring"""
					if not head:
										return True
					# split the list to two parts
					__a      ,			__a					:						Any      =					head.next, head
					while fast and fast.next:
										__a					:						Optional[int]      =					fast.next.next
										__a					:						Optional[int]      =					slow.next
					__a					:						Optional[int]      =					slow.next
					__a					:						Tuple      =					None  # Don't forget here! But forget still works!
					# reverse the second part
					__a					:						Any      =					None
					while second:
										__a					:						int      =					second.next
										__a					:						int      =					node
										__a					:						Union[str, Any]      =					second
										__a					:						Union[str, Any]      =					nxt
					# compare two parts
					# second part has the same or one less node
					while node:
										if node.val != head.val:
															return False
										__a					:						int      =					node.next
										__a					:						Dict      =					head.next
					return True







def     a_   (__A      )				->						Union[str, Any]:



					"""simple docstring"""
					if not head or not head.next:
										return True

					# 1. Get the midpoint (slow)
					__a					:						Union[str, Any]      =					head
					while fast and fast.next:
										__a      ,			__a					:						List[str]      =					fast.next.next, slow.next

					# 2. Push the second half into the stack
					__a					:						Dict      =					[slow.val]
					while slow.next:
										__a					:						Optional[Any]      =					slow.next
										stack.append(slow.val      )

					# 3. Comparison
					while stack:
										if stack.pop() != cur.val:
															return False
										__a					:						List[str]      =					cur.next

					return True







def     a_   (__A      )				->						List[str]:



					"""simple docstring"""
					if not head or not head.next:
										return True
					__a					:						Optional[int]      =					{}
					__a					:						Optional[int]      =					0
					while head:
										if head.val in d:
															d[head.val].append(__A      )
										else:
															__a					:						Dict      =					[pos]
										__a					:						Dict      =					head.next
										pos += 1
					__a					:						Dict      =					pos - 1
					__a					:						Optional[int]      =					0
					for v in d.values():
										if len(__A      ) % 2 != 0:
															middle += 1
										else:
															__a					:						Dict      =					0
															for i in range(0				,       len(__A      )      ):
																				if v[i] + v[len(__A      ) - 1 - step] != checksum:
																									return False
																				step += 1
										if middle > 1:
															return False
					return True





