






'''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 					a_					(					self      : Any					):


            """simple docstring"""







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

                 __lowerCamelCase						:       str				   =				[t[-1] for t in os.walk(os.path.join(A__	,    os.listdir(A__					)[0]	,    """snapshots"""					)					)]
                 __lowerCamelCase						:       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 					a_					(					self      : Union[str, Any]					):


            """simple docstring"""







            __lowerCamelCase   ,      __lowerCamelCase						:       Any				   =				FlaxStableDiffusionPipeline.from_pretrained(
                """hf-internal-testing/tiny-stable-diffusion-pipe"""	,    safety_checker=A__					)

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

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

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

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

            __lowerCamelCase						:       int				   =				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

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




       def 					a_					(					self      : int					):


            """simple docstring"""







            __lowerCamelCase   ,      __lowerCamelCase						:       Optional[Any]				   =				FlaxStableDiffusionPipeline.from_pretrained(
                """CompVis/stable-diffusion-v1-4"""	,    revision="""flax"""	,    safety_checker=A__					)

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

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

            __lowerCamelCase						:       List[Any]				   =				jax.device_count()
            __lowerCamelCase						:       str				   =				num_samples * [prompt]
            __lowerCamelCase						:       Dict				   =				pipeline.prepare_inputs(A__					)

            # shard inputs and rng
            __lowerCamelCase						:       Tuple				   =				replicate(A__					)
            __lowerCamelCase						:       str				   =				jax.random.split(A__	,    A__					)
            __lowerCamelCase						:       Dict				   =				shard(A__					)

            __lowerCamelCase						:       Any				   =				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      : Any					):


            """simple docstring"""







            __lowerCamelCase   ,      __lowerCamelCase						:       Dict				   =				FlaxStableDiffusionPipeline.from_pretrained(
                """CompVis/stable-diffusion-v1-4"""	,    revision="""bf16"""	,    dtype=jnp.bfloataa	,    safety_checker=A__					)

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

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

            __lowerCamelCase						:       Union[str, Any]				   =				jax.device_count()
            __lowerCamelCase						:       int				   =				num_samples * [prompt]
            __lowerCamelCase						:       Optional[int]				   =				pipeline.prepare_inputs(A__					)

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

            __lowerCamelCase						:       int				   =				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      : Dict					):


            """simple docstring"""







            __lowerCamelCase   ,      __lowerCamelCase						:       Any				   =				FlaxStableDiffusionPipeline.from_pretrained(
                """CompVis/stable-diffusion-v1-4"""	,    revision="""bf16"""	,    dtype=jnp.bfloataa					)

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

            __lowerCamelCase						:       str				   =				jax.random.PRNGKey(0					)
            __lowerCamelCase						:       Tuple				   =				50

            __lowerCamelCase						:       Any				   =				jax.device_count()
            __lowerCamelCase						:       Tuple				   =				num_samples * [prompt]
            __lowerCamelCase						:       List[str]				   =				pipeline.prepare_inputs(A__					)

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

            __lowerCamelCase						:       Any				   =				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      : int					):


            """simple docstring"""







            __lowerCamelCase						:       int				   =				FlaxDDIMScheduler(
                beta_start=0.0_0085	,    beta_end=0.012	,    beta_schedule="""scaled_linear"""	,    set_alpha_to_one=A__	,    steps_offset=1	,    )

            __lowerCamelCase   ,      __lowerCamelCase						:       Optional[int]				   =				FlaxStableDiffusionPipeline.from_pretrained(
                """CompVis/stable-diffusion-v1-4"""	,    revision="""bf16"""	,    dtype=jnp.bfloataa	,    scheduler=A__	,    safety_checker=A__	,    )
            __lowerCamelCase						:       List[str]				   =				scheduler.create_state()

            __lowerCamelCase						:       Dict				   =				scheduler_state

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

            __lowerCamelCase						:       Any				   =				jax.random.PRNGKey(0					)
            __lowerCamelCase						:       int				   =				50

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

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

            __lowerCamelCase						:       int				   =				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      : Optional[int]					):


            """simple docstring"""







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

            __lowerCamelCase						:       Optional[int]				   =				jax.device_count()
            __lowerCamelCase						:       Tuple				   =				num_samples * [prompt]
            __lowerCamelCase						:       Optional[Any]				   =				jax.random.split(jax.random.PRNGKey(0					)	,    A__					)

            __lowerCamelCase   ,      __lowerCamelCase						:       Dict				   =				FlaxStableDiffusionPipeline.from_pretrained(
                """CompVis/stable-diffusion-v1-4"""	,    revision="""bf16"""	,    dtype=jnp.bfloataa	,    safety_checker=A__	,    )

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

            # With memory efficient attention
            __lowerCamelCase   ,      __lowerCamelCase						:       Optional[int]				   =				FlaxStableDiffusionPipeline.from_pretrained(
                """CompVis/stable-diffusion-v1-4"""	,    revision="""bf16"""	,    dtype=jnp.bfloataa	,    safety_checker=A__	,    use_memory_efficient_attention=A__	,    )

            __lowerCamelCase						:       int				   =				replicate(A__					)
            __lowerCamelCase						:       Optional[int]				   =				pipeline.prepare_inputs(A__					)
            __lowerCamelCase						:       Union[str, Any]				   =				shard(A__					)
            __lowerCamelCase						:       str				   =				pipeline(A__	,    A__	,    A__	,    jit=A__					).images
            assert images_eff.shape == (num_samples, 1, 512, 512, 3)
            __lowerCamelCase						:       List[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'''

import copy
from dataclasses import dataclass, field
from typing import ClassVar, Dict

from ..features import Audio, Features, Value
from .base import TaskTemplate







@dataclass(frozen=lowerCAmelCase_  )
class 	SCREAMING_SNAKE_CASE	(      lowerCAmelCase_  ):
       snake_case__    :     str   		=			field(default='automatic-speech-recognition'    ,   metadata={'include_in_asdict_even_if_is_default': True}  )
       snake_case__    :     ClassVar[Features]   		=			Features({'audio': Audio()}  )
       snake_case__    :     ClassVar[Features]   		=			Features({'transcription': Value('string'  )}  )
       snake_case__    :     str   		=			"audio"
       snake_case__    :     str   		=			"transcription"




       def 					a_					(					self      : Any	,    A__      : Any					):


            """simple docstring"""







            if self.audio_column not in features:
                 raise ValueError(f"Column {self.audio_column} is not present in features."					)
            if not isinstance(features[self.audio_column]	,    A__					):
                 raise ValueError(f"Column {self.audio_column} is not an Audio type."					)
            __lowerCamelCase						:       Dict				   =				copy.deepcopy(self					)
            __lowerCamelCase						:       List[Any]				   =				self.input_schema.copy()
            __lowerCamelCase						:       Optional[int]				   =				features[self.audio_column]
            __lowerCamelCase						:       Any				   =				input_schema
            return task_template




       @property
       def 					a_					(					self      : Any					):


            """simple docstring"""







            return {self.audio_column: "audio", self.transcription_column: "transcription"}



