


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



    """simple docstring"""






    def _a (      self   )  ->	Dict:
         with tempfile.TemporaryDirectory() as tmpdirname:
              # pipeline has Flax weights
              __UpperCamelCase   							=FlaxDiffusionPipeline.from_pretrained(
                  'hf-internal-testing/tiny-stable-diffusion-pipe' ,       safety_checker=A_ ,       cache_dir=A_   )

              __UpperCamelCase   							=[t[-1] for t in os.walk(os.path.join(A_ ,       os.listdir(A_   )[0] ,       'snapshots'   )   )]
              __UpperCamelCase   							=[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    UpperCAmelCase__	( unittest.TestCase		):



    """simple docstring"""






    def _a (      self   )  ->	Optional[Any]:
         __UpperCamelCase							,				__UpperCamelCase   							=FlaxStableDiffusionPipeline.from_pretrained(
             'hf-internal-testing/tiny-stable-diffusion-pipe' ,       safety_checker=A_   )

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

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

         __UpperCamelCase   							=jax.device_count()
         __UpperCamelCase   							=num_samples * [prompt]
         __UpperCamelCase   							=pipeline.prepare_inputs(A_   )

         # shard inputs and rng
         __UpperCamelCase   							=replicate(A_   )
         __UpperCamelCase   							=jax.random.split(A_ ,       A_   )
         __UpperCamelCase   							=shard(A_   )

         __UpperCamelCase   							=pipeline(A_ ,       A_ ,       A_ ,       A_ ,       jit=A_   ).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.151_4745   ) < 1E-3
              assert np.abs(np.abs(A_ ,       dtype=np.floataa   ).sum() - 4_9947.875   ) < 5E-1

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






    def _a (      self   )  ->	List[str]:
         __UpperCamelCase							,				__UpperCamelCase   							=FlaxStableDiffusionPipeline.from_pretrained(
             'CompVis/stable-diffusion-v1-4' ,       revision='flax' ,       safety_checker=A_   )

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

         __UpperCamelCase   							=jax.random.PRNGKey(0   )
         __UpperCamelCase   							=50

         __UpperCamelCase   							=jax.device_count()
         __UpperCamelCase   							=num_samples * [prompt]
         __UpperCamelCase   							=pipeline.prepare_inputs(A_   )

         # shard inputs and rng
         __UpperCamelCase   							=replicate(A_   )
         __UpperCamelCase   							=jax.random.split(A_ ,       A_   )
         __UpperCamelCase   							=shard(A_   )

         __UpperCamelCase   							=pipeline(A_ ,       A_ ,       A_ ,       A_ ,       jit=A_   ).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.0565_2401)   ) < 1E-3
              assert np.abs((np.abs(A_ ,       dtype=np.floataa   ).sum() - 238_3808.2)   ) < 5E-1






    def _a (      self   )  ->	List[Any]:
         __UpperCamelCase							,				__UpperCamelCase   							=FlaxStableDiffusionPipeline.from_pretrained(
             'CompVis/stable-diffusion-v1-4' ,       revision='bf16' ,       dtype=jnp.bfloataa ,       safety_checker=A_   )

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

         __UpperCamelCase   							=jax.random.PRNGKey(0   )
         __UpperCamelCase   							=50

         __UpperCamelCase   							=jax.device_count()
         __UpperCamelCase   							=num_samples * [prompt]
         __UpperCamelCase   							=pipeline.prepare_inputs(A_   )

         # shard inputs and rng
         __UpperCamelCase   							=replicate(A_   )
         __UpperCamelCase   							=jax.random.split(A_ ,       A_   )
         __UpperCamelCase   							=shard(A_   )

         __UpperCamelCase   							=pipeline(A_ ,       A_ ,       A_ ,       A_ ,       jit=A_   ).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.0400_3906)   ) < 1E-3
              assert np.abs((np.abs(A_ ,       dtype=np.floataa   ).sum() - 237_3516.75)   ) < 5E-1






    def _a (      self   )  ->	List[str]:
         __UpperCamelCase							,				__UpperCamelCase   							=FlaxStableDiffusionPipeline.from_pretrained(
             'CompVis/stable-diffusion-v1-4' ,       revision='bf16' ,       dtype=jnp.bfloataa   )

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

         __UpperCamelCase   							=jax.random.PRNGKey(0   )
         __UpperCamelCase   							=50

         __UpperCamelCase   							=jax.device_count()
         __UpperCamelCase   							=num_samples * [prompt]
         __UpperCamelCase   							=pipeline.prepare_inputs(A_   )

         # shard inputs and rng
         __UpperCamelCase   							=replicate(A_   )
         __UpperCamelCase   							=jax.random.split(A_ ,       A_   )
         __UpperCamelCase   							=shard(A_   )

         __UpperCamelCase   							=pipeline(A_ ,       A_ ,       A_ ,       A_ ,       jit=A_   ).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.0400_3906)   ) < 1E-3
              assert np.abs((np.abs(A_ ,       dtype=np.floataa   ).sum() - 237_3516.75)   ) < 5E-1






    def _a (      self   )  ->	Tuple:
         __UpperCamelCase   							=FlaxDDIMScheduler(
             beta_start=0.0_0085 ,       beta_end=0.012 ,       beta_schedule='scaled_linear' ,       set_alpha_to_one=A_ ,       steps_offset=1 ,       )

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

         __UpperCamelCase   							=scheduler_state

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

         __UpperCamelCase   							=jax.random.PRNGKey(0   )
         __UpperCamelCase   							=50

         __UpperCamelCase   							=jax.device_count()
         __UpperCamelCase   							=num_samples * [prompt]
         __UpperCamelCase   							=pipeline.prepare_inputs(A_   )

         # shard inputs and rng
         __UpperCamelCase   							=replicate(A_   )
         __UpperCamelCase   							=jax.random.split(A_ ,       A_   )
         __UpperCamelCase   							=shard(A_   )

         __UpperCamelCase   							=pipeline(A_ ,       A_ ,       A_ ,       A_ ,       jit=A_   ).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_4504_3945)   ) < 1E-3
              assert np.abs((np.abs(A_ ,       dtype=np.floataa   ).sum() - 234_7693.5)   ) < 5E-1






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

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

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

         __UpperCamelCase   							=replicate(A_   )
         __UpperCamelCase   							=pipeline.prepare_inputs(A_   )
         __UpperCamelCase   							=shard(A_   )
         __UpperCamelCase   							=pipeline(A_ ,       A_ ,       A_ ,       jit=A_   ).images
         assert images.shape == (num_samples, 1, 512, 512, 3)
         __UpperCamelCase   							=images[2, 0, 256, 10:17, 1]

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

         __UpperCamelCase   							=replicate(A_   )
         __UpperCamelCase   							=pipeline.prepare_inputs(A_   )
         __UpperCamelCase   							=shard(A_   )
         __UpperCamelCase   							=pipeline(A_ ,       A_ ,       A_ ,       jit=A_   ).images
         assert images_eff.shape == (num_samples, 1, 512, 512, 3)
         __UpperCamelCase   							=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 ..utils import DummyObject, requires_backends






class    UpperCAmelCase__	( metaclass=A_		):



    """simple docstring"""






    UpperCAmelCase__						:		Any              =  ["speech"]






    def __init__(      self ,       *A_ ,       **A_   )  ->	Any:
         requires_backends(self ,       ['speech']   )






class    UpperCAmelCase__	( metaclass=A_		):



    """simple docstring"""






    UpperCAmelCase__						:		Union[str, Any]              =  ["speech"]






    def __init__(      self ,       *A_ ,       **A_   )  ->	Union[str, Any]:
         requires_backends(self ,       ['speech']   )





