



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






      """simple docstring"""




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

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






      """simple docstring"""




      def lowerCAmelCase_					(       self							:     Optional[Any]   ):
             a__								= FlaxStableDiffusionPipeline.from_pretrained(
                 "hf-internal-testing/tiny-stable-diffusion-pipe"			,safety_checker=a__   )

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

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

             a__								= 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_99_47.8_75   ) < 5e-1

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




      def lowerCAmelCase_					(       self							:     Union[str, Any]   ):
             a__								= FlaxStableDiffusionPipeline.from_pretrained(
                 "CompVis/stable-diffusion-v1-4"			,revision="flax"			,safety_checker=a__   )

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

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

             a__								= pipeline(a__			,a__			,a__			,a__			,jit=a__   ).images

             assert images.shape == (num_samples, 1, 5_12, 5_12, 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() - 2_38_38_08.2)   ) < 5e-1




      def lowerCAmelCase_					(       self							:     List[Any]   ):
             a__								= FlaxStableDiffusionPipeline.from_pretrained(
                 "CompVis/stable-diffusion-v1-4"			,revision="bf16"			,dtype=jnp.bfloataa			,safety_checker=a__   )

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

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

             a__								= pipeline(a__			,a__			,a__			,a__			,jit=a__   ).images

             assert images.shape == (num_samples, 1, 5_12, 5_12, 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() - 2_37_35_16.75)   ) < 5e-1




      def lowerCAmelCase_					(       self							:     Dict   ):
             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(a__   )

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

             a__								= pipeline(a__			,a__			,a__			,a__			,jit=a__   ).images

             assert images.shape == (num_samples, 1, 5_12, 5_12, 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() - 2_37_35_16.75)   ) < 5e-1




      def lowerCAmelCase_					(       self							:     Tuple   ):
             a__								= FlaxDDIMScheduler(
                 beta_start=0.0_0085			,beta_end=0.012			,beta_schedule="scaled_linear"			,set_alpha_to_one=a__			,steps_offset=1			,)

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

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

             a__								= pipeline(a__			,a__			,a__			,a__			,jit=a__   ).images

             assert images.shape == (num_samples, 1, 5_12, 5_12, 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() - 2_34_76_93.5)   ) < 5e-1




      def lowerCAmelCase_					(       self							:     Optional[Any]   ):
             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   )			,a__   )

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

             a__								= replicate(a__   )
             a__								= pipeline.prepare_inputs(a__   )
             a__								= shard(a__   )
             a__								= pipeline(a__			,a__			,a__			,jit=a__   ).images
             assert images.shape == (num_samples, 1, 5_12, 5_12, 3)
             a__								= images[2, 0, 2_56, 10:17, 1]

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

             a__								= replicate(a__   )
             a__								= pipeline.prepare_inputs(a__   )
             a__								= shard(a__   )
             a__								= pipeline(a__			,a__			,a__			,jit=a__   ).images
             assert images_eff.shape == (num_samples, 1, 5_12, 5_12, 3)
             a__								= images[2, 0, 2_56, 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

from packaging import version

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


if TYPE_CHECKING:
			from ...processing_utils import ProcessorMixin
			from ...utils import TensorType


lowercase__			=		logging.get_logger(__name__)

lowercase__			=		{
    """microsoft/layoutlmv3-base""": """https://huggingface.co/microsoft/layoutlmv3-base/resolve/main/config.json""",
}






class lowerCAmelCase__      (	lowercase		):






			'''simple docstring'''



			lowerCamelCase__ =     """layoutlmv3"""
			def __init__(    self    , lowercase=50265    , lowercase=768    , lowercase=12    , lowercase=12    , lowercase=3072    , lowercase="gelu"    , lowercase=0.1    , lowercase=0.1    , lowercase=512    , lowercase=2    , lowercase=0.02    , lowercase=1E-5    , lowercase=1    , lowercase=0    , lowercase=2    , lowercase=1024    , lowercase=128    , lowercase=128    , lowercase=True    , lowercase=32    , lowercase=128    , lowercase=64    , lowercase=256    , lowercase=True    , lowercase=True    , lowercase=True    , lowercase=224    , lowercase=3    , lowercase=16    , lowercase=None    , **lowercase    , ):
							super().__init__(
							    vocab_size=lowercase    , hidden_size=lowercase    , num_hidden_layers=lowercase    , num_attention_heads=lowercase    , intermediate_size=lowercase    , hidden_act=lowercase    , hidden_dropout_prob=lowercase    , attention_probs_dropout_prob=lowercase    , max_position_embeddings=lowercase    , type_vocab_size=lowercase    , initializer_range=lowercase    , layer_norm_eps=lowercase    , pad_token_id=lowercase    , bos_token_id=lowercase    , eos_token_id=lowercase    , **lowercase    , )
							_lowerCamelCase	:     List[str]       =       max_ad_position_embeddings
							_lowerCamelCase	:     Optional[Any]       =       coordinate_size
							_lowerCamelCase	:     int       =       shape_size
							_lowerCamelCase	:     Optional[Any]       =       has_relative_attention_bias
							_lowerCamelCase	:     Dict       =       rel_pos_bins
							_lowerCamelCase	:     Tuple       =       max_rel_pos
							_lowerCamelCase	:     int       =       has_spatial_attention_bias
							_lowerCamelCase	:     Optional[int]       =       rel_ad_pos_bins
							_lowerCamelCase	:     List[Any]       =       max_rel_ad_pos
							_lowerCamelCase	:     Any       =       text_embed
							_lowerCamelCase	:     List[Any]       =       visual_embed
							_lowerCamelCase	:     str       =       input_size
							_lowerCamelCase	:     Optional[Any]       =       num_channels
							_lowerCamelCase	:     Tuple       =       patch_size
							_lowerCamelCase	:     Optional[int]       =       classifier_dropout






class lowerCAmelCase__      (	lowercase		):






			'''simple docstring'''



			lowerCamelCase__ =     version.parse("""1.12"""		)
			@property
			def    A_       (    self    ):
							# The order of inputs is different for question answering and sequence classification
							if self.task in ["question-answering", "sequence-classification"]:
											return OrderedDict(
											    [
											        ('input_ids', {0: 'batch', 1: 'sequence'}),
											        ('attention_mask', {0: 'batch', 1: 'sequence'}),
											        ('bbox', {0: 'batch', 1: 'sequence'}),
											        ('pixel_values', {0: 'batch', 1: 'num_channels', 2: 'height', 3: 'width'}),
											    ]    )
							else:
											return OrderedDict(
											    [
											        ('input_ids', {0: 'batch', 1: 'sequence'}),
											        ('bbox', {0: 'batch', 1: 'sequence'}),
											        ('attention_mask', {0: 'batch', 1: 'sequence'}),
											        ('pixel_values', {0: 'batch', 1: 'num_channels'}),
											    ]    )
			@property
			def    A_       (    self    ):
							return 1E-5
			@property
			def    A_       (    self    ):
							return 12






			def    A_       (    self    , lowercase    , lowercase = -1    , lowercase = -1    , lowercase = False    , lowercase = None    , lowercase = 3    , lowercase = 40    , lowercase = 40    , ):
							setattr(processor.image_processor    , 'apply_ocr'    , lowercase    )

							# If dynamic axis (-1) we forward with a fixed dimension of 2 samples to avoid optimizations made by ONNX
							_lowerCamelCase	:     Any       =       compute_effective_axis_dimension(
							    lowercase    , fixed_dimension=OnnxConfig.default_fixed_batch    , num_token_to_add=0    )
							# If dynamic axis (-1) we forward with a fixed dimension of 8 tokens to avoid optimizations made by ONNX
							_lowerCamelCase	:     Tuple       =       processor.tokenizer.num_special_tokens_to_add(lowercase    )
							_lowerCamelCase	:     Any       =       compute_effective_axis_dimension(
							    lowercase    , fixed_dimension=OnnxConfig.default_fixed_sequence    , num_token_to_add=lowercase    )
							# Generate dummy inputs according to compute batch and sequence
							_lowerCamelCase	:     List[Any]       =       [[' '.join([processor.tokenizer.unk_token]    ) * seq_length]] * batch_size

							# Generate dummy bounding boxes
							_lowerCamelCase	:     List[str]       =       [[[48, 84, 73, 128]]] * batch_size

							# If dynamic axis (-1) we forward with a fixed dimension of 2 samples to avoid optimizations made by ONNX
							# batch_size = compute_effective_axis_dimension(batch_size, fixed_dimension=OnnxConfig.default_fixed_batch)
							_lowerCamelCase	:     Tuple       =       self._generate_dummy_images(lowercase    , lowercase    , lowercase    , lowercase    )

							_lowerCamelCase	:     str       =       dict(
							    processor(
							        lowercase    , text=lowercase    , boxes=lowercase    , return_tensors=lowercase    , )    )

							return inputs