





"""simple docstring"""



import gc
import unittest

from diffusers import FlaxDPMSolverMultistepScheduler, FlaxStableDiffusionPipeline
from diffusers.utils import is_flax_available, slow
from diffusers.utils.testing_utils import require_flax


if is_flax_available():
       import jax
       import jax.numpy as jnp
       from flax.jax_utils import replicate
       from flax.training.common_utils import shard






@slow
@require_flax
class 						__snake_case   (						unittest.TestCase	):





 """simple docstring"""


 def    UpperCamelCase__(       self    ):



     '''simple docstring'''




     super().tearDown()
     gc.collect()

 def    UpperCamelCase__(       self    ):



     '''simple docstring'''




     __A					,			__A       :		List[Any]							=	FlaxStableDiffusionPipeline.from_pretrained(
         '''stabilityai/stable-diffusion-2'''   ,		revision='''bf16'''   ,		dtype=jnp.bfloataa   ,		)

     __A       :		Tuple							=	'''A painting of a squirrel eating a burger'''
     __A       :		Union[str, Any]							=	jax.device_count()
     __A       :		Dict							=	num_samples * [prompt]
     __A       :		Optional[int]							=	sd_pipe.prepare_inputs(__lowerCamelCase    )

     __A       :		Optional[Any]							=	replicate(__lowerCamelCase    )
     __A       :		Union[str, Any]							=	shard(__lowerCamelCase    )

     __A       :		str							=	jax.random.PRNGKey(0    )
     __A       :		List[Any]							=	jax.random.split(__lowerCamelCase   ,		jax.device_count()    )

     __A       :		Any							=	sd_pipe(__lowerCamelCase   ,		__lowerCamelCase   ,		__lowerCamelCase   ,		num_inference_steps=25   ,		jit=__lowerCamelCase    )[0]
     assert images.shape == (jax.device_count(), 1, 768, 768, 3)

     __A       :		int							=	images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]    )
     __A       :		Optional[int]							=	images[0, 253:256, 253:256, -1]

     __A       :		Tuple							=	jnp.asarray(jax.device_get(image_slice.flatten()    )    )
     __A       :		Optional[int]							=	jnp.array([0.4_2_3_8, 0.4_4_1_4, 0.4_3_9_5, 0.4_4_5_3, 0.4_6_2_9, 0.4_5_9_0, 0.4_5_3_1, 0.4_5_5_0_8, 0.4_5_1_2]    )
     print(F"""output_slice: {output_slice}"""    )
     assert jnp.abs(output_slice - expected_slice    ).max() < 1e-2







 def    UpperCamelCase__(       self    ):



     '''simple docstring'''




     __A       :		Any							=	'''stabilityai/stable-diffusion-2'''
     __A					,			__A       :		Union[str, Any]							=	FlaxDPMSolverMultistepScheduler.from_pretrained(__lowerCamelCase   ,		subfolder='''scheduler'''    )
     __A					,			__A       :		Any							=	FlaxStableDiffusionPipeline.from_pretrained(
         __lowerCamelCase   ,		scheduler=__lowerCamelCase   ,		revision='''bf16'''   ,		dtype=jnp.bfloataa   ,		)
     __A       :		List[str]							=	scheduler_params

     __A       :		Any							=	'''A painting of a squirrel eating a burger'''
     __A       :		Dict							=	jax.device_count()
     __A       :		Optional[Any]							=	num_samples * [prompt]
     __A       :		Any							=	sd_pipe.prepare_inputs(__lowerCamelCase    )

     __A       :		str							=	replicate(__lowerCamelCase    )
     __A       :		Any							=	shard(__lowerCamelCase    )

     __A       :		Dict							=	jax.random.PRNGKey(0    )
     __A       :		str							=	jax.random.split(__lowerCamelCase   ,		jax.device_count()    )

     __A       :		Union[str, Any]							=	sd_pipe(__lowerCamelCase   ,		__lowerCamelCase   ,		__lowerCamelCase   ,		num_inference_steps=25   ,		jit=__lowerCamelCase    )[0]
     assert images.shape == (jax.device_count(), 1, 768, 768, 3)

     __A       :		Union[str, Any]							=	images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]    )
     __A       :		Any							=	images[0, 253:256, 253:256, -1]

     __A       :		int							=	jnp.asarray(jax.device_get(image_slice.flatten()    )    )
     __A       :		List[str]							=	jnp.array([0.4_3_3_6, 0.4_2_9_6_9, 0.4_4_5_3, 0.4_1_9_9, 0.4_2_9_7, 0.4_5_3_1, 0.4_4_3_4, 0.4_4_3_4, 0.4_2_9_7]    )
     print(F"""output_slice: {output_slice}"""    )
     assert jnp.abs(output_slice - expected_slice    ).max() < 1e-2











"""simple docstring"""



from collections import OrderedDict
from typing import Mapping

from packaging import version

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


a_					=     logging.get_logger(__name__)

a_					=     {
    """sail/poolformer_s12""": """https://huggingface.co/sail/poolformer_s12/resolve/main/config.json""",
    # See all PoolFormer models at https://huggingface.co/models?filter=poolformer
}






class 						__snake_case   (						SCREAMING_SNAKE_CASE__	):





 """simple docstring"""


 _lowerCamelCase								=	"""poolformer"""

 def __init__(       self   ,		__lowerCamelCase=3   ,		__lowerCamelCase=16   ,		__lowerCamelCase=16   ,		__lowerCamelCase=3   ,		__lowerCamelCase=4.0   ,		__lowerCamelCase=[2, 2, 6, 2]   ,		__lowerCamelCase=[64, 128, 320, 512]   ,		__lowerCamelCase=[7, 3, 3, 3]   ,		__lowerCamelCase=[4, 2, 2, 2]   ,		__lowerCamelCase=[2, 1, 1, 1]   ,		__lowerCamelCase=4   ,		__lowerCamelCase=0.0   ,		__lowerCamelCase="gelu"   ,		__lowerCamelCase=True   ,		__lowerCamelCase=1e-5   ,		__lowerCamelCase=0.0_2   ,		**__lowerCamelCase   ,		):



     '''simple docstring'''




     __A       :		str							=	num_channels
     __A       :		List[str]							=	patch_size
     __A       :		str							=	stride
     __A       :		Any							=	padding
     __A       :		Any							=	pool_size
     __A       :		Dict							=	hidden_sizes
     __A       :		Optional[Any]							=	mlp_ratio
     __A       :		Any							=	depths
     __A       :		List[str]							=	patch_sizes
     __A       :		Union[str, Any]							=	strides
     __A       :		List[str]							=	num_encoder_blocks
     __A       :		Optional[int]							=	drop_path_rate
     __A       :		Union[str, Any]							=	hidden_act
     __A       :		Optional[Any]							=	use_layer_scale
     __A       :		List[Any]							=	layer_scale_init_value
     __A       :		List[Any]							=	initializer_range
     super().__init__(**__lowerCamelCase    )






class 						__snake_case   (						SCREAMING_SNAKE_CASE__	):





 """simple docstring"""


 _lowerCamelCase								=	version.parse("""1.11"""	)

 @property
 def    UpperCamelCase__(       self    ):



     '''simple docstring'''




     return OrderedDict(
         [
             ('''pixel_values''', {0: '''batch''', 1: '''num_channels''', 2: '''height''', 3: '''width'''}),
         ]    )







 @property
 def    UpperCamelCase__(       self    ):



     '''simple docstring'''




     return 2e-3




