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





   """simple docstring"""




   def   lowerCAmelCase						(	self					:  Optional[Any]							):



    """simple docstring"""



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

     _UpperCamelCase:     Optional[Any]							= [t[-1] for t in os.walk(os.path.join(_lowercase		,  os.listdir(_lowercase							)[0]		,  '''snapshots'''							)							)]
     _UpperCamelCase:     int							= [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    ):





   """simple docstring"""




   def   lowerCAmelCase						(	self					:  Union[str, Any]							):



    """simple docstring"""



    _UpperCamelCase	,							_UpperCamelCase:     int							= FlaxStableDiffusionPipeline.from_pretrained(
        '''hf-internal-testing/tiny-stable-diffusion-pipe'''		,  safety_checker=_lowercase							)

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

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

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

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

    _UpperCamelCase:     List[Any]							= 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.1514745							) < 1E-3
     assert np.abs(np.abs(_lowercase		,  dtype=np.floataa							).sum() - 49947.875							) < 5E-1

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



   def   lowerCAmelCase						(	self					:  int							):



    """simple docstring"""



    _UpperCamelCase	,							_UpperCamelCase:     Any							= FlaxStableDiffusionPipeline.from_pretrained(
        '''CompVis/stable-diffusion-v1-4'''		,  revision='''flax'''		,  safety_checker=_lowercase							)

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

    _UpperCamelCase:     Optional[int]							= jax.random.PRNGKey(0							)
    _UpperCamelCase:     Optional[Any]							= 50

    _UpperCamelCase:     Tuple							= jax.device_count()
    _UpperCamelCase:     Optional[int]							= num_samples * [prompt]
    _UpperCamelCase:     List[str]							= pipeline.prepare_inputs(_lowercase							)

    # shard inputs and rng
    _UpperCamelCase:     Dict							= replicate(_lowercase							)
    _UpperCamelCase:     List[str]							= jax.random.split(_lowercase		,  _lowercase							)
    _UpperCamelCase:     List[Any]							= shard(_lowercase							)

    _UpperCamelCase:     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.05652401)							) < 1E-3
     assert np.abs((np.abs(_lowercase		,  dtype=np.floataa							).sum() - 2383808.2)							) < 5E-1



   def   lowerCAmelCase						(	self					:  Optional[Any]							):



    """simple docstring"""



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

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

    _UpperCamelCase:     Optional[Any]							= jax.random.PRNGKey(0							)
    _UpperCamelCase:     str							= 50

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

    # shard inputs and rng
    _UpperCamelCase:     Optional[Any]							= replicate(_lowercase							)
    _UpperCamelCase:     Optional[Any]							= jax.random.split(_lowercase		,  _lowercase							)
    _UpperCamelCase:     Any							= shard(_lowercase							)

    _UpperCamelCase:     Union[str, 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.04003906)							) < 1E-3
     assert np.abs((np.abs(_lowercase		,  dtype=np.floataa							).sum() - 2373516.75)							) < 5E-1



   def   lowerCAmelCase						(	self					:  Optional[Any]							):



    """simple docstring"""



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

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

    _UpperCamelCase:     Dict							= jax.random.PRNGKey(0							)
    _UpperCamelCase:     str							= 50

    _UpperCamelCase:     List[Any]							= jax.device_count()
    _UpperCamelCase:     str							= num_samples * [prompt]
    _UpperCamelCase:     Optional[int]							= pipeline.prepare_inputs(_lowercase							)

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

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



   def   lowerCAmelCase						(	self					:  str							):



    """simple docstring"""



    _UpperCamelCase:     Union[str, Any]							= FlaxDDIMScheduler(
        beta_start=0.00085		,  beta_end=0.012		,  beta_schedule='''scaled_linear'''		,  set_alpha_to_one=_lowercase		,  steps_offset=1		,  )

    _UpperCamelCase	,							_UpperCamelCase:     Optional[int]							= FlaxStableDiffusionPipeline.from_pretrained(
        '''CompVis/stable-diffusion-v1-4'''		,  revision='''bf16'''		,  dtype=jnp.bfloataa		,  scheduler=_lowercase		,  safety_checker=_lowercase		,  )
    _UpperCamelCase:     Optional[Any]							= scheduler.create_state()

    _UpperCamelCase:     Tuple							= scheduler_state

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

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

    _UpperCamelCase:     Tuple							= jax.device_count()
    _UpperCamelCase:     List[str]							= num_samples * [prompt]
    _UpperCamelCase:     str							= pipeline.prepare_inputs(_lowercase							)

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

    _UpperCamelCase:     Union[str, 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.045043945)							) < 1E-3
     assert np.abs((np.abs(_lowercase		,  dtype=np.floataa							).sum() - 2347693.5)							) < 5E-1



   def   lowerCAmelCase						(	self					:  Optional[Any]							):



    """simple docstring"""



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

    _UpperCamelCase:     List[str]							= jax.device_count()
    _UpperCamelCase:     Dict							= num_samples * [prompt]
    _UpperCamelCase:     List[str]							= jax.random.split(jax.random.PRNGKey(0							)		,  _lowercase							)

    _UpperCamelCase	,							_UpperCamelCase:     Any							= FlaxStableDiffusionPipeline.from_pretrained(
        '''CompVis/stable-diffusion-v1-4'''		,  revision='''bf16'''		,  dtype=jnp.bfloataa		,  safety_checker=_lowercase		,  )

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

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

    _UpperCamelCase:     Union[str, Any]							= replicate(_lowercase							)
    _UpperCamelCase:     int							= pipeline.prepare_inputs(_lowercase							)
    _UpperCamelCase:     List[Any]							= shard(_lowercase							)
    _UpperCamelCase:     Optional[int]							= pipeline(_lowercase		,  _lowercase		,  _lowercase		,  jit=_lowercase							).images
    assert images_eff.shape == (num_samples, 1, 512, 512, 3)
    _UpperCamelCase:     int							= 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
from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


UpperCAmelCase_			    =							{
    '''configuration_blenderbot''': [
        '''BLENDERBOT_PRETRAINED_CONFIG_ARCHIVE_MAP''',
        '''BlenderbotConfig''',
        '''BlenderbotOnnxConfig''',
    ],
    '''tokenization_blenderbot''': ['''BlenderbotTokenizer'''],
}

try:
  if not is_tokenizers_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  UpperCAmelCase_			    =							['''BlenderbotTokenizerFast''']

try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  UpperCAmelCase_			    =							[
      '''BLENDERBOT_PRETRAINED_MODEL_ARCHIVE_LIST''',
      '''BlenderbotForCausalLM''',
      '''BlenderbotForConditionalGeneration''',
      '''BlenderbotModel''',
      '''BlenderbotPreTrainedModel''',
  ]


try:
  if not is_tf_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  UpperCAmelCase_			    =							[
      '''TFBlenderbotForConditionalGeneration''',
      '''TFBlenderbotModel''',
      '''TFBlenderbotPreTrainedModel''',
  ]


try:
  if not is_flax_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  UpperCAmelCase_			    =							[
      '''FlaxBlenderbotForConditionalGeneration''',
      '''FlaxBlenderbotModel''',
      '''FlaxBlenderbotPreTrainedModel''',
  ]


if TYPE_CHECKING:
  from .configuration_blenderbot import (
      BLENDERBOT_PRETRAINED_CONFIG_ARCHIVE_MAP,
      BlenderbotConfig,
      BlenderbotOnnxConfig,
  )
  from .tokenization_blenderbot import BlenderbotTokenizer

  try:
    if not is_tokenizers_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .tokenization_blenderbot_fast import BlenderbotTokenizerFast

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_blenderbot import (
        BLENDERBOT_PRETRAINED_MODEL_ARCHIVE_LIST,
        BlenderbotForCausalLM,
        BlenderbotForConditionalGeneration,
        BlenderbotModel,
        BlenderbotPreTrainedModel,
    )

  try:
    if not is_tf_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_tf_blenderbot import (
        TFBlenderbotForConditionalGeneration,
        TFBlenderbotModel,
        TFBlenderbotPreTrainedModel,
    )

  try:
    if not is_flax_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_flax_blenderbot import (
        FlaxBlenderbotForConditionalGeneration,
        FlaxBlenderbotModel,
        FlaxBlenderbotPreTrainedModel,
    )

else:
  import sys

  UpperCAmelCase_			    =							_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)