




"""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    A_       (    self    ):
							with tempfile.TemporaryDirectory() as tmpdirname:
											# pipeline has Flax weights
											_lowerCamelCase	:     Dict       =       FlaxDiffusionPipeline.from_pretrained(
											    'hf-internal-testing/tiny-stable-diffusion-pipe'    , safety_checker=lowercase    , cache_dir=lowercase    )

											_lowerCamelCase	:     int       =       [t[-1] for t in os.walk(os.path.join(lowercase    , os.listdir(lowercase    )[0]    , 'snapshots'    )    )]
											_lowerCamelCase	:     Tuple       =       [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    A_       (    self    ):
							_lowerCamelCase,					_lowerCamelCase	:     Dict       =       FlaxStableDiffusionPipeline.from_pretrained(
							    'hf-internal-testing/tiny-stable-diffusion-pipe'    , safety_checker=lowercase    )

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

							_lowerCamelCase	:     str       =       jax.random.PRNGKey(0    )
							_lowerCamelCase	:     Tuple       =       4

							_lowerCamelCase	:     Optional[Any]       =       jax.device_count()
							_lowerCamelCase	:     Union[str, Any]       =       num_samples * [prompt]
							_lowerCamelCase	:     List[str]       =       pipeline.prepare_inputs(lowercase    )

							# shard inputs and rng
							_lowerCamelCase	:     Any       =       replicate(lowercase    )
							_lowerCamelCase	:     Dict       =       jax.random.split(lowercase    , lowercase    )
							_lowerCamelCase	:     Dict       =       shard(lowercase    )

							_lowerCamelCase	:     List[str]       =       pipeline(lowercase    , lowercase    , lowercase    , lowercase    , jit=lowercase    ).images

							assert images.shape == (num_samples, 1, 64, 64, 3)
							if jax.device_count() == 8:
											assert np.abs(np.abs(images[0, 0, :2, :2, -2:]    , dtype=np.floataa    ).sum() - 4.1_51_47_45    ) < 1E-3
											assert np.abs(np.abs(lowercase    , dtype=np.floataa    ).sum() - 4_99_47.8_75    ) < 5E-1

							_lowerCamelCase	:     Dict       =       pipeline.numpy_to_pil(np.asarray(images.reshape((num_samples,) + images.shape[-3:]    )    )    )
							assert len(lowercase    ) == num_samples
			def    A_       (    self    ):
							_lowerCamelCase,					_lowerCamelCase	:     Union[str, Any]       =       FlaxStableDiffusionPipeline.from_pretrained(
							    'CompVis/stable-diffusion-v1-4'    , revision='flax'    , safety_checker=lowercase    )

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

							_lowerCamelCase	:     Any       =       jax.random.PRNGKey(0    )
							_lowerCamelCase	:     Any       =       50

							_lowerCamelCase	:     Union[str, Any]       =       jax.device_count()
							_lowerCamelCase	:     List[Any]       =       num_samples * [prompt]
							_lowerCamelCase	:     List[Any]       =       pipeline.prepare_inputs(lowercase    )

							# shard inputs and rng
							_lowerCamelCase	:     Optional[Any]       =       replicate(lowercase    )
							_lowerCamelCase	:     Optional[int]       =       jax.random.split(lowercase    , lowercase    )
							_lowerCamelCase	:     Dict       =       shard(lowercase    )

							_lowerCamelCase	:     str       =       pipeline(lowercase    , lowercase    , lowercase    , lowercase    , jit=lowercase    ).images

							assert images.shape == (num_samples, 1, 512, 512, 3)
							if jax.device_count() == 8:
											assert np.abs((np.abs(images[0, 0, :2, :2, -2:]    , dtype=np.floataa    ).sum() - 0.05_65_24_01)    ) < 1E-3
											assert np.abs((np.abs(lowercase    , dtype=np.floataa    ).sum() - 2_38_38_08.2)    ) < 5E-1
			def    A_       (    self    ):
							_lowerCamelCase,					_lowerCamelCase	:     str       =       FlaxStableDiffusionPipeline.from_pretrained(
							    'CompVis/stable-diffusion-v1-4'    , revision='bf16'    , dtype=jnp.bfloataa    , safety_checker=lowercase    )

							_lowerCamelCase	:     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'
							)

							_lowerCamelCase	:     Union[str, Any]       =       jax.random.PRNGKey(0    )
							_lowerCamelCase	:     List[str]       =       50

							_lowerCamelCase	:     Tuple       =       jax.device_count()
							_lowerCamelCase	:     Optional[Any]       =       num_samples * [prompt]
							_lowerCamelCase	:     Any       =       pipeline.prepare_inputs(lowercase    )

							# shard inputs and rng
							_lowerCamelCase	:     Optional[int]       =       replicate(lowercase    )
							_lowerCamelCase	:     Dict       =       jax.random.split(lowercase    , lowercase    )
							_lowerCamelCase	:     List[Any]       =       shard(lowercase    )

							_lowerCamelCase	:     Optional[Any]       =       pipeline(lowercase    , lowercase    , lowercase    , lowercase    , jit=lowercase    ).images

							assert images.shape == (num_samples, 1, 512, 512, 3)
							if jax.device_count() == 8:
											assert np.abs((np.abs(images[0, 0, :2, :2, -2:]    , dtype=np.floataa    ).sum() - 0.04_00_39_06)    ) < 1E-3
											assert np.abs((np.abs(lowercase    , dtype=np.floataa    ).sum() - 2_37_35_16.75)    ) < 5E-1
			def    A_       (    self    ):
							_lowerCamelCase,					_lowerCamelCase	:     List[str]       =       FlaxStableDiffusionPipeline.from_pretrained(
							    'CompVis/stable-diffusion-v1-4'    , revision='bf16'    , dtype=jnp.bfloataa    )

							_lowerCamelCase	:     List[str]       =       (
							    'A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of'
							    ' field, close up, split lighting, cinematic'
							)

							_lowerCamelCase	:     Optional[int]       =       jax.random.PRNGKey(0    )
							_lowerCamelCase	:     Dict       =       50

							_lowerCamelCase	:     Any       =       jax.device_count()
							_lowerCamelCase	:     Tuple       =       num_samples * [prompt]
							_lowerCamelCase	:     Optional[int]       =       pipeline.prepare_inputs(lowercase    )

							# shard inputs and rng
							_lowerCamelCase	:     Any       =       replicate(lowercase    )
							_lowerCamelCase	:     Union[str, Any]       =       jax.random.split(lowercase    , lowercase    )
							_lowerCamelCase	:     str       =       shard(lowercase    )

							_lowerCamelCase	:     str       =       pipeline(lowercase    , lowercase    , lowercase    , lowercase    , jit=lowercase    ).images

							assert images.shape == (num_samples, 1, 512, 512, 3)
							if jax.device_count() == 8:
											assert np.abs((np.abs(images[0, 0, :2, :2, -2:]    , dtype=np.floataa    ).sum() - 0.04_00_39_06)    ) < 1E-3
											assert np.abs((np.abs(lowercase    , dtype=np.floataa    ).sum() - 2_37_35_16.75)    ) < 5E-1
			def    A_       (    self    ):
							_lowerCamelCase	:     str       =       FlaxDDIMScheduler(
							    beta_start=0.0_00_85    , beta_end=0.0_12    , beta_schedule='scaled_linear'    , set_alpha_to_one=lowercase    , steps_offset=1    , )

							_lowerCamelCase,					_lowerCamelCase	:     List[str]       =       FlaxStableDiffusionPipeline.from_pretrained(
							    'CompVis/stable-diffusion-v1-4'    , revision='bf16'    , dtype=jnp.bfloataa    , scheduler=lowercase    , safety_checker=lowercase    , )
							_lowerCamelCase	:     Any       =       scheduler.create_state()

							_lowerCamelCase	:     Any       =       scheduler_state

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

							_lowerCamelCase	:     List[Any]       =       jax.random.PRNGKey(0    )
							_lowerCamelCase	:     Union[str, Any]       =       50

							_lowerCamelCase	:     int       =       jax.device_count()
							_lowerCamelCase	:     str       =       num_samples * [prompt]
							_lowerCamelCase	:     List[Any]       =       pipeline.prepare_inputs(lowercase    )

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

							_lowerCamelCase	:     int       =       pipeline(lowercase    , lowercase    , lowercase    , lowercase    , jit=lowercase    ).images

							assert images.shape == (num_samples, 1, 512, 512, 3)
							if jax.device_count() == 8:
											assert np.abs((np.abs(images[0, 0, :2, :2, -2:]    , dtype=np.floataa    ).sum() - 0.0_45_04_39_45)    ) < 1E-3
											assert np.abs((np.abs(lowercase    , dtype=np.floataa    ).sum() - 2_34_76_93.5)    ) < 5E-1






			def    A_       (    self    ):
							_lowerCamelCase	:     List[str]       =       (
							    'A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of'
							    ' field, close up, split lighting, cinematic'
							)

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

							_lowerCamelCase,					_lowerCamelCase	:     Optional[Any]       =       FlaxStableDiffusionPipeline.from_pretrained(
							    'CompVis/stable-diffusion-v1-4'    , revision='bf16'    , dtype=jnp.bfloataa    , safety_checker=lowercase    , )

							_lowerCamelCase	:     Dict       =       replicate(lowercase    )
							_lowerCamelCase	:     Optional[Any]       =       pipeline.prepare_inputs(lowercase    )
							_lowerCamelCase	:     List[str]       =       shard(lowercase    )
							_lowerCamelCase	:     List[str]       =       pipeline(lowercase    , lowercase    , lowercase    , jit=lowercase    ).images
							assert images.shape == (num_samples, 1, 512, 512, 3)
							_lowerCamelCase	:     Union[str, Any]       =       images[2, 0, 256, 10:17, 1]

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

							_lowerCamelCase	:     Any       =       replicate(lowercase    )
							_lowerCamelCase	:     Any       =       pipeline.prepare_inputs(lowercase    )
							_lowerCamelCase	:     Optional[int]       =       shard(lowercase    )
							_lowerCamelCase	:     List[Any]       =       pipeline(lowercase    , lowercase    , lowercase    , jit=lowercase    ).images
							assert images_eff.shape == (num_samples, 1, 512, 512, 3)
							_lowerCamelCase	:     str       =       images[2, 0, 256, 10:17, 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





"""simple docstring"""






from __future__ import annotations







def 	_snake_case     (							lowercase__						):
				_lowerCamelCase	:     int       =       len(lowercase__						)
				# We need to create solution object to save path.
				_lowerCamelCase	:     Tuple       =       [[0 for _ in range(lowercase__						)] for _ in range(lowercase__						)]
				_lowerCamelCase	:     Optional[Any]       =       run_maze(lowercase__				,       0				,       0				,       lowercase__						)
				if solved:
								print('\n'.join(str(lowercase__						) for row in solutions						)						)
				else:
								print('No solution exists!'						)
				return solved







def 	_snake_case     (							lowercase__				,       lowercase__				,       lowercase__				,       lowercase__						):
				_lowerCamelCase	:     Optional[int]       =       len(lowercase__						)
				# Final check point.
				if i == j == (size - 1):
								_lowerCamelCase	:     Optional[Any]       =       1
								return True

				_lowerCamelCase	:     List[str]       =       (not i < 0) and (not j < 0)  # Check lower bounds
				_lowerCamelCase	:     List[str]       =       (i < size) and (j < size)  # Check upper bounds

				if lower_flag and upper_flag:
								# check for already visited and block points.
								_lowerCamelCase	:     List[Any]       =       (not solutions[i][j]) and (not maze[i][j])
								if block_flag:
												# check visited
												_lowerCamelCase	:     Union[str, Any]       =       1

												# check for directions
												if (
												    run_maze(lowercase__				,       i + 1				,       lowercase__				,       lowercase__						)
												    or run_maze(lowercase__				,       lowercase__				,       j + 1				,       lowercase__						)
												    or run_maze(lowercase__				,       i - 1				,       lowercase__				,       lowercase__						)
												    or run_maze(lowercase__				,       lowercase__				,       j - 1				,       lowercase__						)
												):
																return True

												_lowerCamelCase	:     Dict       =       0
												return False
				return False


if __name__ == "__main__":
			import doctest

			doctest.testmod()