


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 	snake_case_  (	self    ):

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

           __a					=  [t[-1] for t in os.walk(os.path.join(__UpperCAmelCase  ,       os.listdir(__UpperCAmelCase    )[0]  ,       """snapshots"""    )    )]
           __a					=  [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 	snake_case_  (	self    ):
        __a					=  FlaxStableDiffusionPipeline.from_pretrained(
            """hf-internal-testing/tiny-stable-diffusion-pipe"""  ,       safety_checker=__UpperCAmelCase    )

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

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

        __a					=  jax.device_count()
        __a					=  num_samples * [prompt]
        __a					=  pipeline.prepare_inputs(__UpperCAmelCase    )

        # shard inputs and rng
        __a					=  replicate(__UpperCAmelCase    )
        __a					=  jax.random.split(__UpperCAmelCase  ,       __UpperCAmelCase    )
        __a					=  shard(__UpperCAmelCase    )

        __a					=  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					=  pipeline.numpy_to_pil(np.asarray(images.reshape((num_samples,) + images.shape[-3:]    )    )    )
        assert len(__UpperCAmelCase    ) == num_samples


     def 	snake_case_  (	self    ):
        __a					=  FlaxStableDiffusionPipeline.from_pretrained(
            """CompVis/stable-diffusion-v1-4"""  ,       revision="""flax"""  ,       safety_checker=__UpperCAmelCase    )

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

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

        __a					=  jax.device_count()
        __a					=  num_samples * [prompt]
        __a					=  pipeline.prepare_inputs(__UpperCAmelCase    )

        # shard inputs and rng
        __a					=  replicate(__UpperCAmelCase    )
        __a					=  jax.random.split(__UpperCAmelCase  ,       __UpperCAmelCase    )
        __a					=  shard(__UpperCAmelCase    )

        __a					=  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 	snake_case_  (	self    ):
        __a					=  FlaxStableDiffusionPipeline.from_pretrained(
            """CompVis/stable-diffusion-v1-4"""  ,       revision="""bf16"""  ,       dtype=jnp.bfloataa  ,       safety_checker=__UpperCAmelCase    )

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

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

        __a					=  jax.device_count()
        __a					=  num_samples * [prompt]
        __a					=  pipeline.prepare_inputs(__UpperCAmelCase    )

        # shard inputs and rng
        __a					=  replicate(__UpperCAmelCase    )
        __a					=  jax.random.split(__UpperCAmelCase  ,       __UpperCAmelCase    )
        __a					=  shard(__UpperCAmelCase    )

        __a					=  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 	snake_case_  (	self    ):
        __a					=  FlaxStableDiffusionPipeline.from_pretrained(
            """CompVis/stable-diffusion-v1-4"""  ,       revision="""bf16"""  ,       dtype=jnp.bfloataa    )

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

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

        __a					=  jax.device_count()
        __a					=  num_samples * [prompt]
        __a					=  pipeline.prepare_inputs(__UpperCAmelCase    )

        # shard inputs and rng
        __a					=  replicate(__UpperCAmelCase    )
        __a					=  jax.random.split(__UpperCAmelCase  ,       __UpperCAmelCase    )
        __a					=  shard(__UpperCAmelCase    )

        __a					=  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 	snake_case_  (	self    ):
        __a					=  FlaxDDIMScheduler(
            beta_start=0.00085  ,       beta_end=0.012  ,       beta_schedule="""scaled_linear"""  ,       set_alpha_to_one=__UpperCAmelCase  ,       steps_offset=1  ,       )

        __a					=  FlaxStableDiffusionPipeline.from_pretrained(
            """CompVis/stable-diffusion-v1-4"""  ,       revision="""bf16"""  ,       dtype=jnp.bfloataa  ,       scheduler=__UpperCAmelCase  ,       safety_checker=__UpperCAmelCase  ,       )
        __a					=  scheduler.create_state()

        __a					=  scheduler_state

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

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

        __a					=  jax.device_count()
        __a					=  num_samples * [prompt]
        __a					=  pipeline.prepare_inputs(__UpperCAmelCase    )

        # shard inputs and rng
        __a					=  replicate(__UpperCAmelCase    )
        __a					=  jax.random.split(__UpperCAmelCase  ,       __UpperCAmelCase    )
        __a					=  shard(__UpperCAmelCase    )

        __a					=  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 	snake_case_  (	self    ):
        __a					=  (
            'A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of'
            ' field, close up, split lighting, cinematic'
        )

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

        __a					=  FlaxStableDiffusionPipeline.from_pretrained(
            """CompVis/stable-diffusion-v1-4"""  ,       revision="""bf16"""  ,       dtype=jnp.bfloataa  ,       safety_checker=__UpperCAmelCase  ,       )

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

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

        __a					=  replicate(__UpperCAmelCase    )
        __a					=  pipeline.prepare_inputs(__UpperCAmelCase    )
        __a					=  shard(__UpperCAmelCase    )
        __a					=  pipeline(__UpperCAmelCase  ,       __UpperCAmelCase  ,       __UpperCAmelCase  ,       jit=__UpperCAmelCase    ).images
        assert images_eff.shape == (num_samples, 1, 512, 512, 3)
        __a					=  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 collections import OrderedDict
from typing import TYPE_CHECKING, Any, Mapping, Optional, Union

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


if TYPE_CHECKING:
  from ... import FeatureExtractionMixin, PreTrainedTokenizerBase, TensorType


__A											=						logging.get_logger(__name__)

__A											=						{
    """microsoft/deberta-v2-xlarge""": """https://huggingface.co/microsoft/deberta-v2-xlarge/resolve/main/config.json""",
    """microsoft/deberta-v2-xxlarge""": """https://huggingface.co/microsoft/deberta-v2-xxlarge/resolve/main/config.json""",
    """microsoft/deberta-v2-xlarge-mnli""": (
        """https://huggingface.co/microsoft/deberta-v2-xlarge-mnli/resolve/main/config.json"""
    ),
    """microsoft/deberta-v2-xxlarge-mnli""": (
        """https://huggingface.co/microsoft/deberta-v2-xxlarge-mnli/resolve/main/config.json"""
    ),
}






class    _lowerCAmelCase  (	a			):
       """simple docstring"""




       __magic_name__	:Union[str, Any]						    =   """deberta-v2"""






       def __init__(       self       ,					__UpperCAmelCase=1_2_8_1_0_0       ,					__UpperCAmelCase=1_5_3_6       ,					__UpperCAmelCase=2_4       ,					__UpperCAmelCase=2_4       ,					__UpperCAmelCase=6_1_4_4       ,					__UpperCAmelCase="gelu"       ,					__UpperCAmelCase=0.1       ,					__UpperCAmelCase=0.1       ,					__UpperCAmelCase=5_1_2       ,					__UpperCAmelCase=0       ,					__UpperCAmelCase=0.02       ,					__UpperCAmelCase=1E-7       ,					__UpperCAmelCase=False       ,					__UpperCAmelCase=-1       ,					__UpperCAmelCase=0       ,					__UpperCAmelCase=True       ,					__UpperCAmelCase=None       ,					__UpperCAmelCase=0       ,					__UpperCAmelCase="gelu"       ,					**__UpperCAmelCase       ,					):



         '''simple docstring'''

         super().__init__(**__UpperCAmelCase )

         lowerCAmelCase__     :Optional[int]											=				hidden_size
         lowerCAmelCase__     :Optional[int]											=				num_hidden_layers
         lowerCAmelCase__     :Optional[int]											=				num_attention_heads
         lowerCAmelCase__     :int											=				intermediate_size
         lowerCAmelCase__     :Dict											=				hidden_act
         lowerCAmelCase__     :Tuple											=				hidden_dropout_prob
         lowerCAmelCase__     :Optional[int]											=				attention_probs_dropout_prob
         lowerCAmelCase__     :Optional[Any]											=				max_position_embeddings
         lowerCAmelCase__     :Optional[Any]											=				type_vocab_size
         lowerCAmelCase__     :List[Any]											=				initializer_range
         lowerCAmelCase__     :Dict											=				relative_attention
         lowerCAmelCase__     :Union[str, Any]											=				max_relative_positions
         lowerCAmelCase__     :List[str]											=				pad_token_id
         lowerCAmelCase__     :Optional[int]											=				position_biased_input

         # Backwards compatibility
         if type(__UpperCAmelCase ) == str:
           lowerCAmelCase__     :Dict											=				[x.strip() for x in pos_att_type.lower().split('|' )]

         lowerCAmelCase__     :str											=				pos_att_type
         lowerCAmelCase__     :Dict											=				vocab_size
         lowerCAmelCase__     :Tuple											=				layer_norm_eps

         lowerCAmelCase__     :Optional[Any]											=				kwargs.get('pooler_hidden_size'       ,					__UpperCAmelCase )
         lowerCAmelCase__     :Union[str, Any]											=				pooler_dropout
         lowerCAmelCase__     :int											=				pooler_hidden_act






class    _lowerCAmelCase  (	a			):
       """simple docstring"""




       @property
       def 	snake_case		(       self ):



         '''simple docstring'''

         if self.task == "multiple-choice":
           lowerCAmelCase__     :str											=				{0: 'batch', 1: 'choice', 2: 'sequence'}
         else:
           lowerCAmelCase__     :Tuple											=				{0: 'batch', 1: 'sequence'}
         if self._config.type_vocab_size > 0:
           return OrderedDict(
               [('input_ids', dynamic_axis), ('attention_mask', dynamic_axis), ('token_type_ids', dynamic_axis)] )
         else:
           return OrderedDict([('input_ids', dynamic_axis), ('attention_mask', dynamic_axis)] )






       @property
       def 	snake_case		(       self ):



         '''simple docstring'''

         return 1_2






       def 	snake_case		(       self       ,					__UpperCAmelCase       ,					__UpperCAmelCase = -1       ,					__UpperCAmelCase = -1       ,					__UpperCAmelCase = -1       ,					__UpperCAmelCase = False       ,					__UpperCAmelCase = None       ,					__UpperCAmelCase = 3       ,					__UpperCAmelCase = 4_0       ,					__UpperCAmelCase = 4_0       ,					__UpperCAmelCase = None       ,					):



         '''simple docstring'''

         lowerCAmelCase__     :Union[str, Any]											=				super().generate_dummy_inputs(preprocessor=__UpperCAmelCase       ,					framework=__UpperCAmelCase )
         if self._config.type_vocab_size == 0 and "token_type_ids" in dummy_inputs:
           del dummy_inputs["token_type_ids"]
         return dummy_inputs
