



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

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

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

									_lowerCamelCase        =				jax.random.PRNGKey(0   )
									_lowerCamelCase        =				4

									_lowerCamelCase        =				jax.device_count()
									_lowerCamelCase        =				num_samples * [prompt]
									_lowerCamelCase        =				pipeline.prepare_inputs(__lowercase   )

									# shard inputs and rng
									_lowerCamelCase        =				replicate(__lowercase   )
									_lowerCamelCase        =				jax.random.split(__lowercase				,       __lowercase   )
									_lowerCamelCase        =				shard(__lowercase   )

									_lowerCamelCase        =				pipeline(__lowercase				,       __lowercase				,       __lowercase				,       __lowercase				,       jit=__lowercase   ).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.1_5_1_4_7_4_5   ) < 1e-3
													assert np.abs(np.abs(__lowercase				,       dtype=np.floataa   ).sum() - 4_9_9_4_7.8_7_5   ) < 5e-1

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

					def 			snake_case__      (    self   ):
									_lowerCamelCase  ,     _lowerCamelCase        =				FlaxStableDiffusionPipeline.from_pretrained(
									    '''CompVis/stable-diffusion-v1-4'''				,       revision='''flax'''				,       safety_checker=__lowercase   )

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

									_lowerCamelCase        =				jax.random.PRNGKey(0   )
									_lowerCamelCase        =				5_0

									_lowerCamelCase        =				jax.device_count()
									_lowerCamelCase        =				num_samples * [prompt]
									_lowerCamelCase        =				pipeline.prepare_inputs(__lowercase   )

									# shard inputs and rng
									_lowerCamelCase        =				replicate(__lowercase   )
									_lowerCamelCase        =				jax.random.split(__lowercase				,       __lowercase   )
									_lowerCamelCase        =				shard(__lowercase   )

									_lowerCamelCase        =				pipeline(__lowercase				,       __lowercase				,       __lowercase				,       __lowercase				,       jit=__lowercase   ).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_5_6_5_2_4_0_1)   ) < 1e-3
													assert np.abs((np.abs(__lowercase				,       dtype=np.floataa   ).sum() - 2_3_8_3_8_0_8.2)   ) < 5e-1

					def 			snake_case__      (    self   ):
									_lowerCamelCase  ,     _lowerCamelCase        =				FlaxStableDiffusionPipeline.from_pretrained(
									    '''CompVis/stable-diffusion-v1-4'''				,       revision='''bf16'''				,       dtype=jnp.bfloataa				,       safety_checker=__lowercase   )

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

									_lowerCamelCase        =				jax.random.PRNGKey(0   )
									_lowerCamelCase        =				5_0

									_lowerCamelCase        =				jax.device_count()
									_lowerCamelCase        =				num_samples * [prompt]
									_lowerCamelCase        =				pipeline.prepare_inputs(__lowercase   )

									# shard inputs and rng
									_lowerCamelCase        =				replicate(__lowercase   )
									_lowerCamelCase        =				jax.random.split(__lowercase				,       __lowercase   )
									_lowerCamelCase        =				shard(__lowercase   )

									_lowerCamelCase        =				pipeline(__lowercase				,       __lowercase				,       __lowercase				,       __lowercase				,       jit=__lowercase   ).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_4_0_0_3_9_0_6)   ) < 1e-3
													assert np.abs((np.abs(__lowercase				,       dtype=np.floataa   ).sum() - 2_3_7_3_5_1_6.7_5)   ) < 5e-1

					def 			snake_case__      (    self   ):
									_lowerCamelCase  ,     _lowerCamelCase        =				FlaxStableDiffusionPipeline.from_pretrained(
									    '''CompVis/stable-diffusion-v1-4'''				,       revision='''bf16'''				,       dtype=jnp.bfloataa   )

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

									_lowerCamelCase        =				jax.random.PRNGKey(0   )
									_lowerCamelCase        =				5_0

									_lowerCamelCase        =				jax.device_count()
									_lowerCamelCase        =				num_samples * [prompt]
									_lowerCamelCase        =				pipeline.prepare_inputs(__lowercase   )

									# shard inputs and rng
									_lowerCamelCase        =				replicate(__lowercase   )
									_lowerCamelCase        =				jax.random.split(__lowercase				,       __lowercase   )
									_lowerCamelCase        =				shard(__lowercase   )

									_lowerCamelCase        =				pipeline(__lowercase				,       __lowercase				,       __lowercase				,       __lowercase				,       jit=__lowercase   ).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_4_0_0_3_9_0_6)   ) < 1e-3
													assert np.abs((np.abs(__lowercase				,       dtype=np.floataa   ).sum() - 2_3_7_3_5_1_6.7_5)   ) < 5e-1

					def 			snake_case__      (    self   ):
									_lowerCamelCase        =				FlaxDDIMScheduler(
									    beta_start=0.0_0_0_8_5				,       beta_end=0.0_1_2				,       beta_schedule='''scaled_linear'''				,       set_alpha_to_one=__lowercase				,       steps_offset=1				,       )

									_lowerCamelCase  ,     _lowerCamelCase        =				FlaxStableDiffusionPipeline.from_pretrained(
									    '''CompVis/stable-diffusion-v1-4'''				,       revision='''bf16'''				,       dtype=jnp.bfloataa				,       scheduler=__lowercase				,       safety_checker=__lowercase				,       )
									_lowerCamelCase        =				scheduler.create_state()

									_lowerCamelCase        =				scheduler_state

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

									_lowerCamelCase        =				jax.random.PRNGKey(0   )
									_lowerCamelCase        =				5_0

									_lowerCamelCase        =				jax.device_count()
									_lowerCamelCase        =				num_samples * [prompt]
									_lowerCamelCase        =				pipeline.prepare_inputs(__lowercase   )

									# shard inputs and rng
									_lowerCamelCase        =				replicate(__lowercase   )
									_lowerCamelCase        =				jax.random.split(__lowercase				,       __lowercase   )
									_lowerCamelCase        =				shard(__lowercase   )

									_lowerCamelCase        =				pipeline(__lowercase				,       __lowercase				,       __lowercase				,       __lowercase				,       jit=__lowercase   ).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_4_5_0_4_3_9_4_5)   ) < 1e-3
													assert np.abs((np.abs(__lowercase				,       dtype=np.floataa   ).sum() - 2_3_4_7_6_9_3.5)   ) < 5e-1





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

									_lowerCamelCase        =				jax.device_count()
									_lowerCamelCase        =				num_samples * [prompt]
									_lowerCamelCase        =				jax.random.split(jax.random.PRNGKey(0   )				,       __lowercase   )

									_lowerCamelCase  ,     _lowerCamelCase        =				FlaxStableDiffusionPipeline.from_pretrained(
									    '''CompVis/stable-diffusion-v1-4'''				,       revision='''bf16'''				,       dtype=jnp.bfloataa				,       safety_checker=__lowercase				,       )

									_lowerCamelCase        =				replicate(__lowercase   )
									_lowerCamelCase        =				pipeline.prepare_inputs(__lowercase   )
									_lowerCamelCase        =				shard(__lowercase   )
									_lowerCamelCase        =				pipeline(__lowercase				,       __lowercase				,       __lowercase				,       jit=__lowercase   ).images
									assert images.shape == (num_samples, 1, 5_1_2, 5_1_2, 3)
									_lowerCamelCase        =				images[2, 0, 2_5_6, 1_0:1_7, 1]

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

									_lowerCamelCase        =				replicate(__lowercase   )
									_lowerCamelCase        =				pipeline.prepare_inputs(__lowercase   )
									_lowerCamelCase        =				shard(__lowercase   )
									_lowerCamelCase        =				pipeline(__lowercase				,       __lowercase				,       __lowercase				,       jit=__lowercase   ).images
									assert images_eff.shape == (num_samples, 1, 5_1_2, 5_1_2, 3)
									_lowerCamelCase        =				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





"""simple docstring"""
import warnings

from .generation import TFGenerationMixin


class 							lowerCamelCase_(   A__    ):





					'''simple docstring'''



					warnings.warn(
					    'Importing `TFGenerationMixin` from `src/transformers/generation_tf_utils.py` is deprecated and will '
					    'be removed in Transformers v5. Import as `from transformers import TFGenerationMixin` instead.',				A__,				)
