



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





						def        _snake_case    (self				):





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

												lowerCamelCase__ :      Union[str, Any]								=		[t[-1] for t in os.walk(os.path.join(UpperCamelCase__							,    os.listdir(UpperCamelCase__				)[0]							,    """snapshots"""				)				)]
												lowerCamelCase__ :      List[Any]								=		[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 				__A					( unittest.TestCase      ):





						def        _snake_case    (self				):
									lowerCamelCase__ :      List[Any]								=		FlaxStableDiffusionPipeline.from_pretrained(
									    """hf-internal-testing/tiny-stable-diffusion-pipe"""							,    safety_checker=UpperCamelCase__				)

									lowerCamelCase__ :      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.random.PRNGKey(0				)
									lowerCamelCase__ :      Optional[int]								=		4

									lowerCamelCase__ :      Optional[int]								=		jax.device_count()
									lowerCamelCase__ :      Dict								=		num_samples * [prompt]
									lowerCamelCase__ :      Any								=		pipeline.prepare_inputs(UpperCamelCase__				)

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

									lowerCamelCase__ :      List[str]								=		pipeline(UpperCamelCase__							,    UpperCamelCase__							,    UpperCamelCase__							,    UpperCamelCase__							,    jit=UpperCamelCase__				).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(UpperCamelCase__							,    dtype=np.floataa				).sum() - 49947.875				) < 5E-1

									lowerCamelCase__ :      str								=		pipeline.numpy_to_pil(np.asarray(images.reshape((num_samples,) + images.shape[-3:]				)				)				)
									assert len(UpperCamelCase__				) == num_samples





						def        _snake_case    (self				):
									lowerCamelCase__ :      Optional[Any]								=		FlaxStableDiffusionPipeline.from_pretrained(
									    """CompVis/stable-diffusion-v1-4"""							,    revision="""flax"""							,    safety_checker=UpperCamelCase__				)

									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[str]								=		jax.random.PRNGKey(0				)
									lowerCamelCase__ :      Optional[int]								=		50

									lowerCamelCase__ :      str								=		jax.device_count()
									lowerCamelCase__ :      Union[str, Any]								=		num_samples * [prompt]
									lowerCamelCase__ :      Optional[Any]								=		pipeline.prepare_inputs(UpperCamelCase__				)

									# shard inputs and rng
									lowerCamelCase__ :      Optional[Any]								=		replicate(UpperCamelCase__				)
									lowerCamelCase__ :      Dict								=		jax.random.split(UpperCamelCase__							,    UpperCamelCase__				)
									lowerCamelCase__ :      Union[str, Any]								=		shard(UpperCamelCase__				)

									lowerCamelCase__ :      Dict								=		pipeline(UpperCamelCase__							,    UpperCamelCase__							,    UpperCamelCase__							,    UpperCamelCase__							,    jit=UpperCamelCase__				).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(UpperCamelCase__							,    dtype=np.floataa				).sum() - 2383808.2)				) < 5E-1





						def        _snake_case    (self				):
									lowerCamelCase__ :      Union[str, Any]								=		FlaxStableDiffusionPipeline.from_pretrained(
									    """CompVis/stable-diffusion-v1-4"""							,    revision="""bf16"""							,    dtype=jnp.bfloataa							,    safety_checker=UpperCamelCase__				)

									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__ :      List[str]								=		jax.random.PRNGKey(0				)
									lowerCamelCase__ :      Union[str, Any]								=		50

									lowerCamelCase__ :      Optional[int]								=		jax.device_count()
									lowerCamelCase__ :      int								=		num_samples * [prompt]
									lowerCamelCase__ :      List[Any]								=		pipeline.prepare_inputs(UpperCamelCase__				)

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

									lowerCamelCase__ :      Tuple								=		pipeline(UpperCamelCase__							,    UpperCamelCase__							,    UpperCamelCase__							,    UpperCamelCase__							,    jit=UpperCamelCase__				).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(UpperCamelCase__							,    dtype=np.floataa				).sum() - 2373516.75)				) < 5E-1





						def        _snake_case    (self				):
									lowerCamelCase__ :      Dict								=		FlaxStableDiffusionPipeline.from_pretrained(
									    """CompVis/stable-diffusion-v1-4"""							,    revision="""bf16"""							,    dtype=jnp.bfloataa				)

									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__ :      Optional[Any]								=		jax.random.PRNGKey(0				)
									lowerCamelCase__ :      Optional[Any]								=		50

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

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

									lowerCamelCase__ :      Optional[int]								=		pipeline(UpperCamelCase__							,    UpperCamelCase__							,    UpperCamelCase__							,    UpperCamelCase__							,    jit=UpperCamelCase__				).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(UpperCamelCase__							,    dtype=np.floataa				).sum() - 2373516.75)				) < 5E-1





						def        _snake_case    (self				):
									lowerCamelCase__ :      Tuple								=		FlaxDDIMScheduler(
									    beta_start=0.0_00_85							,    beta_end=0.0_12							,    beta_schedule="""scaled_linear"""							,    set_alpha_to_one=UpperCamelCase__							,    steps_offset=1							,    )

									lowerCamelCase__ :      Any								=		FlaxStableDiffusionPipeline.from_pretrained(
									    """CompVis/stable-diffusion-v1-4"""							,    revision="""bf16"""							,    dtype=jnp.bfloataa							,    scheduler=UpperCamelCase__							,    safety_checker=UpperCamelCase__							,    )
									lowerCamelCase__ :      Any								=		scheduler.create_state()

									lowerCamelCase__ :      Tuple								=		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__ :      Optional[Any]								=		jax.random.PRNGKey(0				)
									lowerCamelCase__ :      str								=		50

									lowerCamelCase__ :      Tuple								=		jax.device_count()
									lowerCamelCase__ :      Optional[Any]								=		num_samples * [prompt]
									lowerCamelCase__ :      str								=		pipeline.prepare_inputs(UpperCamelCase__				)

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

									lowerCamelCase__ :      Optional[Any]								=		pipeline(UpperCamelCase__							,    UpperCamelCase__							,    UpperCamelCase__							,    UpperCamelCase__							,    jit=UpperCamelCase__				).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(UpperCamelCase__							,    dtype=np.floataa				).sum() - 2347693.5)				) < 5E-1





						def        _snake_case    (self				):
									lowerCamelCase__ :      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__ :      int								=		jax.device_count()
									lowerCamelCase__ :      Any								=		num_samples * [prompt]
									lowerCamelCase__ :      List[str]								=		jax.random.split(jax.random.PRNGKey(0				)							,    UpperCamelCase__				)

									lowerCamelCase__ :      Dict								=		FlaxStableDiffusionPipeline.from_pretrained(
									    """CompVis/stable-diffusion-v1-4"""							,    revision="""bf16"""							,    dtype=jnp.bfloataa							,    safety_checker=UpperCamelCase__							,    )

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

									# With memory efficient attention
									lowerCamelCase__ :      Tuple								=		FlaxStableDiffusionPipeline.from_pretrained(
									    """CompVis/stable-diffusion-v1-4"""							,    revision="""bf16"""							,    dtype=jnp.bfloataa							,    safety_checker=UpperCamelCase__							,    use_memory_efficient_attention=UpperCamelCase__							,    )

									lowerCamelCase__ :      str								=		replicate(UpperCamelCase__				)
									lowerCamelCase__ :      int								=		pipeline.prepare_inputs(UpperCamelCase__				)
									lowerCamelCase__ :      Optional[Any]								=		shard(UpperCamelCase__				)
									lowerCamelCase__ :      Optional[int]								=		pipeline(UpperCamelCase__							,    UpperCamelCase__							,    UpperCamelCase__							,    jit=UpperCamelCase__				).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








import unittest

from .lib import (
    Matrix,
    Vector,
    axpy,
    square_zero_matrix,
    unit_basis_vector,
    zero_vector,
)
class 				__A					( unittest.TestCase      ):





			def        _snake_case    (self				):
						lowerCamelCase__ :      List[str]								=		Vector([1, 2, 3]				)
						self.assertEqual(x.component(0				)							,    1				)
						self.assertEqual(x.component(2				)							,    3				)
						lowerCamelCase__ :      str								=		Vector()





			def        _snake_case    (self				):
						lowerCamelCase__ :      Dict								=		Vector([0, 0, 0, 0, 0, 1]				)
						self.assertEqual(str(__magic_name__				)							,    """(0,0,0,0,0,1)"""				)





			def        _snake_case    (self				):
						lowerCamelCase__ :      Tuple								=		Vector([1, 2, 3, 4]				)
						self.assertEqual(len(__magic_name__				)							,    4				)





			def        _snake_case    (self				):
						lowerCamelCase__ :      Dict								=		Vector([1, 2]				)
						lowerCamelCase__ :      str								=		Vector([1, 2, 3, 4, 5]				)
						lowerCamelCase__ :      int								=		Vector([0, 0, 0, 0, 0, 0, 0, 0, 0, 0]				)
						lowerCamelCase__ :      int								=		Vector([1, -1, 1, -1, 2, -3, 4, -5]				)
						self.assertAlmostEqual(x.euclidean_length()							,    2.2_36							,    3				)
						self.assertAlmostEqual(y.euclidean_length()							,    7.4_16							,    3				)
						self.assertEqual(z.euclidean_length()							,    0				)
						self.assertAlmostEqual(w.euclidean_length()							,    7.6_16							,    3				)





			def        _snake_case    (self				):
						lowerCamelCase__ :      str								=		Vector([1, 2, 3]				)
						lowerCamelCase__ :      Union[str, Any]								=		Vector([1, 1, 1]				)
						self.assertEqual((x + y).component(0				)							,    2				)
						self.assertEqual((x + y).component(1				)							,    3				)
						self.assertEqual((x + y).component(2				)							,    4				)





			def        _snake_case    (self				):
						lowerCamelCase__ :      Optional[int]								=		Vector([1, 2, 3]				)
						lowerCamelCase__ :      Any								=		Vector([1, 1, 1]				)
						self.assertEqual((x - y).component(0				)							,    0				)
						self.assertEqual((x - y).component(1				)							,    1				)
						self.assertEqual((x - y).component(2				)							,    2				)





			def        _snake_case    (self				):
						lowerCamelCase__ :      Dict								=		Vector([1, 2, 3]				)
						lowerCamelCase__ :      int								=		Vector([2, -1, 4]				)  # for test of dot product
						lowerCamelCase__ :      List[Any]								=		Vector([1, -2, -1]				)
						self.assertEqual(str(x * 3.0				)							,    """(3.0,6.0,9.0)"""				)
						self.assertEqual((a * b)							,    0				)





			def        _snake_case    (self				):
						self.assertEqual(str(zero_vector(10				)				).count("""0"""				)							,    10				)





			def        _snake_case    (self				):
						self.assertEqual(str(unit_basis_vector(3							,    1				)				)							,    """(0,1,0)"""				)





			def        _snake_case    (self				):
						lowerCamelCase__ :      Any								=		Vector([1, 2, 3]				)
						lowerCamelCase__ :      str								=		Vector([1, 0, 1]				)
						self.assertEqual(str(axpy(2							,    __magic_name__							,    __magic_name__				)				)							,    """(3,4,7)"""				)





			def        _snake_case    (self				):
						lowerCamelCase__ :      Optional[Any]								=		Vector([1, 0, 0, 0, 0, 0]				)
						lowerCamelCase__ :      int								=		x.copy()
						self.assertEqual(str(__magic_name__				)							,    str(__magic_name__				)				)





			def        _snake_case    (self				):
						lowerCamelCase__ :      int								=		Vector([1, 0, 0]				)
						x.change_component(0							,    0				)
						x.change_component(1							,    1				)
						self.assertEqual(str(__magic_name__				)							,    """(0,1,0)"""				)





			def        _snake_case    (self				):
						lowerCamelCase__ :      Optional[Any]								=		Matrix([[1, 2, 3], [2, 4, 5], [6, 7, 8]]							,    3							,    3				)
						self.assertEqual("""|1,2,3|\n|2,4,5|\n|6,7,8|\n"""							,    str(__magic_name__				)				)





			def        _snake_case    (self				):
						lowerCamelCase__ :      int								=		Matrix([[1, 2, 3], [2, 4, 5], [6, 7, 8]]							,    3							,    3				)
						lowerCamelCase__ :      int								=		[[-3, -14, -10], [-5, -10, -5], [-2, -1, 0]]
						for x in range(a.height()				):
									for y in range(a.width()				):
												self.assertEqual(minors[x][y]							,    a.minor(__magic_name__							,    __magic_name__				)				)





			def        _snake_case    (self				):
						lowerCamelCase__ :      Any								=		Matrix([[1, 2, 3], [2, 4, 5], [6, 7, 8]]							,    3							,    3				)
						lowerCamelCase__ :      Optional[Any]								=		[[-3, 14, -10], [5, -10, 5], [-2, 1, 0]]
						for x in range(a.height()				):
									for y in range(a.width()				):
												self.assertEqual(cofactors[x][y]							,    a.cofactor(__magic_name__							,    __magic_name__				)				)





			def        _snake_case    (self				):
						lowerCamelCase__ :      Optional[int]								=		Matrix([[1, 2, 3], [2, 4, 5], [6, 7, 8]]							,    3							,    3				)
						self.assertEqual(-5							,    a.determinant()				)





			def        _snake_case    (self				):
						lowerCamelCase__ :      List[str]								=		Matrix([[1, 2, 3], [4, 5, 6], [7, 8, 9]]							,    3							,    3				)
						lowerCamelCase__ :      Optional[int]								=		Vector([1, 2, 3]				)
						self.assertEqual("""(14,32,50)"""							,    str(a * x				)				)
						self.assertEqual("""|2,4,6|\n|8,10,12|\n|14,16,18|\n"""							,    str(a * 2				)				)





			def        _snake_case    (self				):
						lowerCamelCase__ :      str								=		Matrix([[1, 2, 3], [2, 4, 5], [6, 7, 8]]							,    3							,    3				)
						a.change_component(0							,    2							,    5				)
						self.assertEqual("""|1,2,5|\n|2,4,5|\n|6,7,8|\n"""							,    str(__magic_name__				)				)





			def        _snake_case    (self				):
						lowerCamelCase__ :      List[Any]								=		Matrix([[1, 2, 3], [2, 4, 5], [6, 7, 8]]							,    3							,    3				)
						self.assertEqual(7							,    a.component(2							,    1				)							,    0.01				)





			def        _snake_case    (self				):
						lowerCamelCase__ :      Tuple								=		Matrix([[1, 2, 3], [2, 4, 5], [6, 7, 8]]							,    3							,    3				)
						lowerCamelCase__ :      int								=		Matrix([[1, 2, 7], [2, 4, 5], [6, 7, 10]]							,    3							,    3				)
						self.assertEqual("""|2,4,10|\n|4,8,10|\n|12,14,18|\n"""							,    str(a + b				)				)





			def        _snake_case    (self				):
						lowerCamelCase__ :      str								=		Matrix([[1, 2, 3], [2, 4, 5], [6, 7, 8]]							,    3							,    3				)
						lowerCamelCase__ :      Optional[Any]								=		Matrix([[1, 2, 7], [2, 4, 5], [6, 7, 10]]							,    3							,    3				)
						self.assertEqual("""|0,0,-4|\n|0,0,0|\n|0,0,-2|\n"""							,    str(a - b				)				)





			def        _snake_case    (self				):
						self.assertEqual(
						    """|0,0,0,0,0|\n|0,0,0,0,0|\n|0,0,0,0,0|\n|0,0,0,0,0|\n|0,0,0,0,0|\n"""							,    str(square_zero_matrix(5				)				)							,    )


if __name__ == "__main__":
				unittest.main()



