





"""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 __magic_name__   (       unittest.TestCase		):
       def 		_lowerCamelCase   (     self	):





         """simple docstring"""


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

           _lowerCAmelCase	      =					[t[-1] for t in os.walk(os.path.join(__lowerCAmelCase   ,  os.listdir(__lowerCAmelCase	)[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 __magic_name__   (       unittest.TestCase		):
       def 		_lowerCamelCase   (     self	):





         """simple docstring"""


         _lowerCAmelCase	      =					FlaxStableDiffusionPipeline.from_pretrained(
             'hf-internal-testing/tiny-stable-diffusion-pipe'   ,  safety_checker=__lowerCAmelCase	)

         _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(__lowerCAmelCase	)

         # shard inputs and rng
         _lowerCAmelCase	      =					replicate(__lowerCAmelCase	)
         _lowerCAmelCase	      =					jax.random.split(__lowerCAmelCase   ,  __lowerCAmelCase	)
         _lowerCAmelCase	      =					shard(__lowerCAmelCase	)

         _lowerCAmelCase	      =					pipeline(__lowerCAmelCase   ,  __lowerCAmelCase   ,  __lowerCAmelCase   ,  __lowerCAmelCase   ,  jit=__lowerCAmelCase	).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_51_47_45	) < 1e-3
           assert np.abs(np.abs(__lowerCAmelCase   ,  dtype=np.floataa	).sum() - 4_99_47.8_75	) < 5e-1

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





         """simple docstring"""


         _lowerCAmelCase	      =					FlaxStableDiffusionPipeline.from_pretrained(
             'CompVis/stable-diffusion-v1-4'   ,  revision='flax'   ,  safety_checker=__lowerCAmelCase	)

         _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(__lowerCAmelCase	)

         # shard inputs and rng
         _lowerCAmelCase	      =					replicate(__lowerCAmelCase	)
         _lowerCAmelCase	      =					jax.random.split(__lowerCAmelCase   ,  __lowerCAmelCase	)
         _lowerCAmelCase	      =					shard(__lowerCAmelCase	)

         _lowerCAmelCase	      =					pipeline(__lowerCAmelCase   ,  __lowerCAmelCase   ,  __lowerCAmelCase   ,  __lowerCAmelCase   ,  jit=__lowerCAmelCase	).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.05_65_24_01)	) < 1e-3
           assert np.abs((np.abs(__lowerCAmelCase   ,  dtype=np.floataa	).sum() - 2_38_38_08.2)	) < 5e-1
       def 		_lowerCamelCase   (     self	):





         """simple docstring"""


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

         _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(__lowerCAmelCase	)

         # shard inputs and rng
         _lowerCAmelCase	      =					replicate(__lowerCAmelCase	)
         _lowerCAmelCase	      =					jax.random.split(__lowerCAmelCase   ,  __lowerCAmelCase	)
         _lowerCAmelCase	      =					shard(__lowerCAmelCase	)

         _lowerCAmelCase	      =					pipeline(__lowerCAmelCase   ,  __lowerCAmelCase   ,  __lowerCAmelCase   ,  __lowerCAmelCase   ,  jit=__lowerCAmelCase	).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.04_00_39_06)	) < 1e-3
           assert np.abs((np.abs(__lowerCAmelCase   ,  dtype=np.floataa	).sum() - 2_37_35_16.75)	) < 5e-1
       def 		_lowerCamelCase   (     self	):





         """simple docstring"""


         _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(__lowerCAmelCase	)

         # shard inputs and rng
         _lowerCAmelCase	      =					replicate(__lowerCAmelCase	)
         _lowerCAmelCase	      =					jax.random.split(__lowerCAmelCase   ,  __lowerCAmelCase	)
         _lowerCAmelCase	      =					shard(__lowerCAmelCase	)

         _lowerCAmelCase	      =					pipeline(__lowerCAmelCase   ,  __lowerCAmelCase   ,  __lowerCAmelCase   ,  __lowerCAmelCase   ,  jit=__lowerCAmelCase	).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.04_00_39_06)	) < 1e-3
           assert np.abs((np.abs(__lowerCAmelCase   ,  dtype=np.floataa	).sum() - 2_37_35_16.75)	) < 5e-1
       def 		_lowerCamelCase   (     self	):





         """simple docstring"""


         _lowerCAmelCase	      =					FlaxDDIMScheduler(
             beta_start=0.0_00_85   ,  beta_end=0.0_12   ,  beta_schedule='scaled_linear'   ,  set_alpha_to_one=__lowerCAmelCase   ,  steps_offset=1   ,  )

         _lowerCAmelCase	      =					FlaxStableDiffusionPipeline.from_pretrained(
             'CompVis/stable-diffusion-v1-4'   ,  revision='bf16'   ,  dtype=jnp.bfloataa   ,  scheduler=__lowerCAmelCase   ,  safety_checker=__lowerCAmelCase   ,  )
         _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(__lowerCAmelCase	)

         # shard inputs and rng
         _lowerCAmelCase	      =					replicate(__lowerCAmelCase	)
         _lowerCAmelCase	      =					jax.random.split(__lowerCAmelCase   ,  __lowerCAmelCase	)
         _lowerCAmelCase	      =					shard(__lowerCAmelCase	)

         _lowerCAmelCase	      =					pipeline(__lowerCAmelCase   ,  __lowerCAmelCase   ,  __lowerCAmelCase   ,  __lowerCAmelCase   ,  jit=__lowerCAmelCase	).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_45_04_39_45)	) < 1e-3
           assert np.abs((np.abs(__lowerCAmelCase   ,  dtype=np.floataa	).sum() - 2_34_76_93.5)	) < 5e-1
       def 		_lowerCamelCase   (     self	):





         """simple docstring"""


         _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	)   ,  __lowerCAmelCase	)

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

         _lowerCAmelCase	      =					replicate(__lowerCAmelCase	)
         _lowerCAmelCase	      =					pipeline.prepare_inputs(__lowerCAmelCase	)
         _lowerCAmelCase	      =					shard(__lowerCAmelCase	)
         _lowerCAmelCase	      =					pipeline(__lowerCAmelCase   ,  __lowerCAmelCase   ,  __lowerCAmelCase   ,  jit=__lowerCAmelCase	).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	      =					FlaxStableDiffusionPipeline.from_pretrained(
             'CompVis/stable-diffusion-v1-4'   ,  revision='bf16'   ,  dtype=jnp.bfloataa   ,  safety_checker=__lowerCAmelCase   ,  use_memory_efficient_attention=__lowerCAmelCase   ,  )

         _lowerCAmelCase	      =					replicate(__lowerCAmelCase	)
         _lowerCAmelCase	      =					pipeline.prepare_inputs(__lowerCAmelCase	)
         _lowerCAmelCase	      =					shard(__lowerCAmelCase	)
         _lowerCAmelCase	      =					pipeline(__lowerCAmelCase   ,  __lowerCAmelCase   ,  __lowerCAmelCase   ,  jit=__lowerCAmelCase	).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











from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tokenizers_available, is_torch_available


SCREAMING_SNAKE_CASE__   :      Dict				=		{
    """configuration_canine""": ["""CANINE_PRETRAINED_CONFIG_ARCHIVE_MAP""", """CanineConfig"""],
    """tokenization_canine""": ["""CanineTokenizer"""],
}

try:
					if not is_torch_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					SCREAMING_SNAKE_CASE__   :      str				=		[
					    """CANINE_PRETRAINED_MODEL_ARCHIVE_LIST""",
					    """CanineForMultipleChoice""",
					    """CanineForQuestionAnswering""",
					    """CanineForSequenceClassification""",
					    """CanineForTokenClassification""",
					    """CanineLayer""",
					    """CanineModel""",
					    """CaninePreTrainedModel""",
					    """load_tf_weights_in_canine""",
					]


if TYPE_CHECKING:
					from .configuration_canine import CANINE_PRETRAINED_CONFIG_ARCHIVE_MAP, CanineConfig
					from .tokenization_canine import CanineTokenizer

					try:
										if not is_torch_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .modeling_canine import (
										    CANINE_PRETRAINED_MODEL_ARCHIVE_LIST,
										    CanineForMultipleChoice,
										    CanineForQuestionAnswering,
										    CanineForSequenceClassification,
										    CanineForTokenClassification,
										    CanineLayer,
										    CanineModel,
										    CaninePreTrainedModel,
										    load_tf_weights_in_canine,
										)


else:
					import sys

					SCREAMING_SNAKE_CASE__   :      Union[str, Any]				=		_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)






