





from collections import OrderedDict
from typing import Mapping

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


__lowercase  :   List[str]						      =logging.get_logger(__name__)

__lowercase  :   Optional[int]						      ={
    """roberta-base""": """https://huggingface.co/roberta-base/resolve/main/config.json""",
    """roberta-large""": """https://huggingface.co/roberta-large/resolve/main/config.json""",
    """roberta-large-mnli""": """https://huggingface.co/roberta-large-mnli/resolve/main/config.json""",
    """distilroberta-base""": """https://huggingface.co/distilroberta-base/resolve/main/config.json""",
    """roberta-base-openai-detector""": """https://huggingface.co/roberta-base-openai-detector/resolve/main/config.json""",
    """roberta-large-openai-detector""": """https://huggingface.co/roberta-large-openai-detector/resolve/main/config.json""",
}

class      A       (	__lowercase						):
      _snake_case							     ='''roberta'''



      def __init__(	self:						List[str]		,  _lowerCAmelCase:						List[str]=5_0265		,  _lowerCAmelCase:						Union[str, Any]=768		,  _lowerCAmelCase:						List[Any]=12		,  _lowerCAmelCase:						Dict=12		,  _lowerCAmelCase:						List[Any]=3072		,  _lowerCAmelCase:						Dict="gelu"		,  _lowerCAmelCase:						Union[str, Any]=0.1		,  _lowerCAmelCase:						Optional[Any]=0.1		,  _lowerCAmelCase:						List[str]=512		,  _lowerCAmelCase:						Optional[Any]=2		,  _lowerCAmelCase:						Optional[Any]=0.02		,  _lowerCAmelCase:						Union[str, Any]=1e-12		,  _lowerCAmelCase:						str=1		,  _lowerCAmelCase:						Any=0		,  _lowerCAmelCase:						Optional[int]=2		,  _lowerCAmelCase:						int="absolute"		,  _lowerCAmelCase:						Optional[Any]=True		,  _lowerCAmelCase:						Tuple=None		,  **_lowerCAmelCase:						Optional[Any]		,  )						->    Optional[Any]:
         '''simple docstring'''

         super().__init__(pad_token_id=_lowerCAmelCase		,  bos_token_id=_lowerCAmelCase		,  eos_token_id=_lowerCAmelCase		,  **_lowerCAmelCase       )

         UpperCAmelCase_											=vocab_size
         UpperCAmelCase_											=hidden_size
         UpperCAmelCase_											=num_hidden_layers
         UpperCAmelCase_											=num_attention_heads
         UpperCAmelCase_											=hidden_act
         UpperCAmelCase_											=intermediate_size
         UpperCAmelCase_											=hidden_dropout_prob
         UpperCAmelCase_											=attention_probs_dropout_prob
         UpperCAmelCase_											=max_position_embeddings
         UpperCAmelCase_											=type_vocab_size
         UpperCAmelCase_											=initializer_range
         UpperCAmelCase_											=layer_norm_eps
         UpperCAmelCase_											=position_embedding_type
         UpperCAmelCase_											=use_cache
         UpperCAmelCase_											=classifier_dropout

class      A       (	__lowercase						):



      @property
      def 	lowerCAmelCase__    (	self:						Any       )						->    Mapping[str, Mapping[int, str]]:
         '''simple docstring'''

         if self.task == "multiple-choice":
            UpperCAmelCase_											={0: "batch", 1: "choice", 2: "sequence"}
         else:
            UpperCAmelCase_											={0: "batch", 1: "sequence"}
         return OrderedDict(
             [
                 ("input_ids", dynamic_axis),
                 ("attention_mask", dynamic_axis),
             ]       )







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



      def 	lowerCAmelCase__    (	self:						str       )						->    Union[str, Any]:
         '''simple docstring'''

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

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



      def 	lowerCAmelCase__    (	self:						Union[str, Any]       )						->    Tuple:
         '''simple docstring'''

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

         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(_lowerCAmelCase       )

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

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

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



      def 	lowerCAmelCase__    (	self:						List[Any]       )						->    Optional[Any]:
         '''simple docstring'''

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

         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(_lowerCAmelCase       )

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

         UpperCAmelCase_											=pipeline(_lowerCAmelCase		,  _lowerCAmelCase		,  _lowerCAmelCase		,  _lowerCAmelCase		,  jit=_lowerCAmelCase       ).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.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:						int       )						->    Union[str, Any]:
         '''simple docstring'''

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

         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(_lowerCAmelCase       )

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

         UpperCAmelCase_											=pipeline(_lowerCAmelCase		,  _lowerCAmelCase		,  _lowerCAmelCase		,  _lowerCAmelCase		,  jit=_lowerCAmelCase       ).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.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:						Optional[int]       )						->    Tuple:
         '''simple docstring'''

         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(_lowerCAmelCase       )

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

         UpperCAmelCase_											=pipeline(_lowerCAmelCase		,  _lowerCAmelCase		,  _lowerCAmelCase		,  _lowerCAmelCase		,  jit=_lowerCAmelCase       ).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.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:						Optional[Any]       )						->    str:
         '''simple docstring'''

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

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

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

         UpperCAmelCase_											=pipeline(_lowerCAmelCase		,  _lowerCAmelCase		,  _lowerCAmelCase		,  _lowerCAmelCase		,  jit=_lowerCAmelCase       ).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_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:						Optional[Any]       )						->    int:
         '''simple docstring'''

         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       )		,  _lowerCAmelCase       )

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

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

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




