





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


					def 							SCREAMING_SNAKE_CASE  (						self):

										'''simple docstring'''






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

															__A						:			int								=	[t[-1] for t in os.walk(os.path.join(_UpperCAmelCase						,     os.listdir(_UpperCAmelCase)[0]						,     'snapshots'))]
															__A						:			List[str]								=	[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 			SCREAMING_SNAKE_CASE   (unittest.TestCase							):


					def 							SCREAMING_SNAKE_CASE  (						self):

										'''simple docstring'''






										__A						:			str								=	FlaxStableDiffusionPipeline.from_pretrained(
										    'hf-internal-testing/tiny-stable-diffusion-pipe'						,     safety_checker=_UpperCAmelCase)

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

										__A						:			List[Any]								=	jax.random.PRNGKey(0)
										__A						:			Dict								=	4

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

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

										__A						:			Optional[int]								=	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.1514745) < 1e-3
															assert np.abs(np.abs(_UpperCAmelCase						,     dtype=np.floataa).sum() - 49947.875) < 5e-1

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


					def 							SCREAMING_SNAKE_CASE  (						self):

										'''simple docstring'''






										__A						:			int								=	FlaxStableDiffusionPipeline.from_pretrained(
										    'CompVis/stable-diffusion-v1-4'						,     revision='flax'						,     safety_checker=_UpperCAmelCase)

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

										__A						:			Optional[int]								=	jax.random.PRNGKey(0)
										__A						:			str								=	50

										__A						:			Optional[int]								=	jax.device_count()
										__A						:			List[str]								=	num_samples * [prompt]
										__A						:			str								=	pipeline.prepare_inputs(_UpperCAmelCase)

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

										__A						:			List[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.05652401)) < 1e-3
															assert np.abs((np.abs(_UpperCAmelCase						,     dtype=np.floataa).sum() - 2383808.2)) < 5e-1


					def 							SCREAMING_SNAKE_CASE  (						self):

										'''simple docstring'''






										__A						:			List[str]								=	FlaxStableDiffusionPipeline.from_pretrained(
										    'CompVis/stable-diffusion-v1-4'						,     revision='bf16'						,     dtype=jnp.bfloataa						,     safety_checker=_UpperCAmelCase)

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

										__A						:			Tuple								=	jax.random.PRNGKey(0)
										__A						:			List[str]								=	50

										__A						:			Dict								=	jax.device_count()
										__A						:			Dict								=	num_samples * [prompt]
										__A						:			str								=	pipeline.prepare_inputs(_UpperCAmelCase)

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

										__A						:			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.04003906)) < 1e-3
															assert np.abs((np.abs(_UpperCAmelCase						,     dtype=np.floataa).sum() - 2373516.75)) < 5e-1


					def 							SCREAMING_SNAKE_CASE  (						self):

										'''simple docstring'''






										__A						:			Dict								=	FlaxStableDiffusionPipeline.from_pretrained(
										    'CompVis/stable-diffusion-v1-4'						,     revision='bf16'						,     dtype=jnp.bfloataa)

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

										__A						:			str								=	jax.random.PRNGKey(0)
										__A						:			List[Any]								=	50

										__A						:			str								=	jax.device_count()
										__A						:			List[str]								=	num_samples * [prompt]
										__A						:			int								=	pipeline.prepare_inputs(_UpperCAmelCase)

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

										__A						:			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.04003906)) < 1e-3
															assert np.abs((np.abs(_UpperCAmelCase						,     dtype=np.floataa).sum() - 2373516.75)) < 5e-1


					def 							SCREAMING_SNAKE_CASE  (						self):

										'''simple docstring'''






										__A						:			int								=	FlaxDDIMScheduler(
										    beta_start=0.00085						,     beta_end=0.012						,     beta_schedule='scaled_linear'						,     set_alpha_to_one=_UpperCAmelCase						,     steps_offset=1						,     )

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

										__A						:			str								=	scheduler_state

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

										__A						:			Optional[Any]								=	jax.random.PRNGKey(0)
										__A						:			List[Any]								=	50

										__A						:			Tuple								=	jax.device_count()
										__A						:			Optional[int]								=	num_samples * [prompt]
										__A						:			int								=	pipeline.prepare_inputs(_UpperCAmelCase)

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

										__A						:			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.045043945)) < 1e-3
															assert np.abs((np.abs(_UpperCAmelCase						,     dtype=np.floataa).sum() - 2347693.5)) < 5e-1


					def 							SCREAMING_SNAKE_CASE  (						self):

										'''simple docstring'''






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

										__A						:			List[str]								=	jax.device_count()
										__A						:			Optional[Any]								=	num_samples * [prompt]
										__A						:			List[str]								=	jax.random.split(jax.random.PRNGKey(0)						,     _UpperCAmelCase)

										__A						:			str								=	FlaxStableDiffusionPipeline.from_pretrained(
										    'CompVis/stable-diffusion-v1-4'						,     revision='bf16'						,     dtype=jnp.bfloataa						,     safety_checker=_UpperCAmelCase						,     )

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

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

										__A						:			Optional[int]								=	replicate(_UpperCAmelCase)
										__A						:			int								=	pipeline.prepare_inputs(_UpperCAmelCase)
										__A						:			Tuple								=	shard(_UpperCAmelCase)
										__A						:			Tuple								=	pipeline(_UpperCAmelCase						,     _UpperCAmelCase						,     _UpperCAmelCase						,     jit=_UpperCAmelCase).images
										assert images_eff.shape == (num_samples, 1, 512, 512, 3)
										__A						:			Any								=	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 tempfile import TemporaryDirectory
from unittest import TestCase
from unittest.mock import MagicMock, patch

from transformers import AutoModel, TFAutoModel
from transformers.onnx import FeaturesManager
from transformers.testing_utils import SMALL_MODEL_IDENTIFIER, require_tf, require_torch


@require_torch
@require_tf
class 			SCREAMING_SNAKE_CASE   (a__							):


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						:			Optional[int]								=	SMALL_MODEL_IDENTIFIER
          __A						:			Any								=	'pt'
          __A						:			str								=	'tf'


     def 							SCREAMING_SNAKE_CASE  (						self						,     _UpperCAmelCase):

          '''simple docstring'''






          __A						:			Any								=	AutoModel.from_pretrained(self.test_model)
          model_pt.save_pretrained(_UpperCAmelCase)


     def 							SCREAMING_SNAKE_CASE  (						self						,     _UpperCAmelCase):

          '''simple docstring'''






          __A						:			Union[str, Any]								=	TFAutoModel.from_pretrained(self.test_model						,     from_pt=_UpperCAmelCase)
          model_tf.save_pretrained(_UpperCAmelCase)


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						:			Union[str, Any]								=	'mock_framework'

          # Framework provided - return whatever the user provides
          __A						:			Union[str, Any]								=	FeaturesManager.determine_framework(self.test_model						,     _UpperCAmelCase)
          self.assertEqual(_UpperCAmelCase						,     _UpperCAmelCase)

          # Local checkpoint and framework provided - return provided framework
          # PyTorch checkpoint
          with TemporaryDirectory() as local_pt_ckpt:
               self._setup_pt_ckpt(_UpperCAmelCase)
               __A						:			List[str]								=	FeaturesManager.determine_framework(_UpperCAmelCase						,     _UpperCAmelCase)
               self.assertEqual(_UpperCAmelCase						,     _UpperCAmelCase)

          # TensorFlow checkpoint
          with TemporaryDirectory() as local_tf_ckpt:
               self._setup_tf_ckpt(_UpperCAmelCase)
               __A						:			Tuple								=	FeaturesManager.determine_framework(_UpperCAmelCase						,     _UpperCAmelCase)
               self.assertEqual(_UpperCAmelCase						,     _UpperCAmelCase)


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          with TemporaryDirectory() as local_pt_ckpt:
               self._setup_pt_ckpt(_UpperCAmelCase)
               __A						:			Optional[int]								=	FeaturesManager.determine_framework(_UpperCAmelCase)
               self.assertEqual(_UpperCAmelCase						,     self.framework_pt)

          # TensorFlow checkpoint
          with TemporaryDirectory() as local_tf_ckpt:
               self._setup_tf_ckpt(_UpperCAmelCase)
               __A						:			List[str]								=	FeaturesManager.determine_framework(_UpperCAmelCase)
               self.assertEqual(_UpperCAmelCase						,     self.framework_tf)

          # Invalid local checkpoint
          with TemporaryDirectory() as local_invalid_ckpt:
               with self.assertRaises(_UpperCAmelCase):
                    __A						:			Tuple								=	FeaturesManager.determine_framework(_UpperCAmelCase)


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						:			Tuple								=	MagicMock(return_value=_UpperCAmelCase)
          with patch('transformers.onnx.features.is_tf_available'						,     _UpperCAmelCase):
               __A						:			Union[str, Any]								=	FeaturesManager.determine_framework(self.test_model)
               self.assertEqual(_UpperCAmelCase						,     self.framework_pt)

          # PyTorch not in environment -> use TensorFlow
          __A						:			List[str]								=	MagicMock(return_value=_UpperCAmelCase)
          with patch('transformers.onnx.features.is_torch_available'						,     _UpperCAmelCase):
               __A						:			List[Any]								=	FeaturesManager.determine_framework(self.test_model)
               self.assertEqual(_UpperCAmelCase						,     self.framework_tf)

          # Both in environment -> use PyTorch
          __A						:			Any								=	MagicMock(return_value=_UpperCAmelCase)
          __A						:			Dict								=	MagicMock(return_value=_UpperCAmelCase)
          with patch('transformers.onnx.features.is_tf_available'						,     _UpperCAmelCase), patch(
              'transformers.onnx.features.is_torch_available'						,     _UpperCAmelCase):
               __A						:			int								=	FeaturesManager.determine_framework(self.test_model)
               self.assertEqual(_UpperCAmelCase						,     self.framework_pt)

          # Both not in environment -> raise error
          __A						:			List[str]								=	MagicMock(return_value=_UpperCAmelCase)
          __A						:			Tuple								=	MagicMock(return_value=_UpperCAmelCase)
          with patch('transformers.onnx.features.is_tf_available'						,     _UpperCAmelCase), patch(
              'transformers.onnx.features.is_torch_available'						,     _UpperCAmelCase):
               with self.assertRaises(_UpperCAmelCase):
                    __A						:			int								=	FeaturesManager.determine_framework(self.test_model)