





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


						def     snake_case_   (self		) ->   Optional[Any]:
								# clean up the VRAM after each test
								super().tearDown()
								gc.collect()


						def     snake_case_   (self		) ->   List[Any]:
								UpperCamelCase   ,				UpperCamelCase   =	FlaxStableDiffusionPipeline.from_pretrained(
								    "stabilityai/stable-diffusion-2"   ,							revision="bf16"   ,							dtype=jnp.bfloataa   ,							)

								UpperCamelCase   =	"A painting of a squirrel eating a burger"
								UpperCamelCase   =	jax.device_count()
								UpperCamelCase   =	num_samples * [prompt]
								UpperCamelCase   =	sd_pipe.prepare_inputs(_lowerCAmelCase		)

								UpperCamelCase   =	replicate(_lowerCAmelCase		)
								UpperCamelCase   =	shard(_lowerCAmelCase		)

								UpperCamelCase   =	jax.random.PRNGKey(0		)
								UpperCamelCase   =	jax.random.split(_lowerCAmelCase   ,							jax.device_count()		)

								UpperCamelCase   =	sd_pipe(_lowerCAmelCase   ,							_lowerCAmelCase   ,							_lowerCAmelCase   ,							num_inference_steps=25   ,							jit=_lowerCAmelCase		)[0]
								assert images.shape == (jax.device_count(), 1, 7_68, 7_68, 3)

								UpperCamelCase   =	images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]		)
								UpperCamelCase   =	images[0, 2_53:2_56, 2_53:2_56, -1]

								UpperCamelCase   =	jnp.asarray(jax.device_get(image_slice.flatten()		)		)
								UpperCamelCase   =	jnp.array([0.4238, 0.4414, 0.4395, 0.4453, 0.4629, 0.4590, 0.4531, 0.45508, 0.4512]		)
								print(F"output_slice: {output_slice}"		)
								assert jnp.abs(output_slice - expected_slice		).max() < 1e-2




						def     snake_case_   (self		) ->   Tuple:
								UpperCamelCase   =	"stabilityai/stable-diffusion-2"
								UpperCamelCase   ,				UpperCamelCase   =	FlaxDPMSolverMultistepScheduler.from_pretrained(_lowerCAmelCase   ,							subfolder="scheduler"		)
								UpperCamelCase   ,				UpperCamelCase   =	FlaxStableDiffusionPipeline.from_pretrained(
								    _lowerCAmelCase   ,							scheduler=_lowerCAmelCase   ,							revision="bf16"   ,							dtype=jnp.bfloataa   ,							)
								UpperCamelCase   =	scheduler_params

								UpperCamelCase   =	"A painting of a squirrel eating a burger"
								UpperCamelCase   =	jax.device_count()
								UpperCamelCase   =	num_samples * [prompt]
								UpperCamelCase   =	sd_pipe.prepare_inputs(_lowerCAmelCase		)

								UpperCamelCase   =	replicate(_lowerCAmelCase		)
								UpperCamelCase   =	shard(_lowerCAmelCase		)

								UpperCamelCase   =	jax.random.PRNGKey(0		)
								UpperCamelCase   =	jax.random.split(_lowerCAmelCase   ,							jax.device_count()		)

								UpperCamelCase   =	sd_pipe(_lowerCAmelCase   ,							_lowerCAmelCase   ,							_lowerCAmelCase   ,							num_inference_steps=25   ,							jit=_lowerCAmelCase		)[0]
								assert images.shape == (jax.device_count(), 1, 7_68, 7_68, 3)

								UpperCamelCase   =	images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]		)
								UpperCamelCase   =	images[0, 2_53:2_56, 2_53:2_56, -1]

								UpperCamelCase   =	jnp.asarray(jax.device_get(image_slice.flatten()		)		)
								UpperCamelCase   =	jnp.array([0.4336, 0.42969, 0.4453, 0.4199, 0.4297, 0.4531, 0.4434, 0.4434, 0.4297]		)
								print(F"output_slice: {output_slice}"		)
								assert jnp.abs(output_slice - expected_slice		).max() < 1e-2













"""simple docstring"""





from ...configuration_utils import PretrainedConfig
from ...utils import logging


lowerCAmelCase__     			=			logging.get_logger(__name__)

lowerCAmelCase__     			=			{
    '''microsoft/swinv2-tiny-patch4-window8-256''': (
        '''https://huggingface.co/microsoft/swinv2-tiny-patch4-window8-256/resolve/main/config.json'''
    ),
}





class        _lowerCamelCase      (       _lowercase ):
						UpperCAmelCase_ =			"swinv2"

						UpperCAmelCase_ =			{
						    "num_attention_heads": "num_heads",
						    "num_hidden_layers": "num_layers",
						}


						def __init__(self   ,							__a=2_24   ,							__a=4   ,							__a=3   ,							__a=96   ,							__a=[2, 2, 6, 2]   ,							__a=[3, 6, 12, 24]   ,							__a=7   ,							__a=4.0   ,							__a=True   ,							__a=0.0   ,							__a=0.0   ,							__a=0.1   ,							__a="gelu"   ,							__a=False   ,							__a=0.02   ,							__a=1e-5   ,							__a=32   ,							**__a   ,							) ->   Optional[Any]:
								super().__init__(**__a		)

								UpperCamelCase   =	image_size
								UpperCamelCase   =	patch_size
								UpperCamelCase   =	num_channels
								UpperCamelCase   =	embed_dim
								UpperCamelCase   =	depths
								UpperCamelCase   =	len(__a		)
								UpperCamelCase   =	num_heads
								UpperCamelCase   =	window_size
								UpperCamelCase   =	mlp_ratio
								UpperCamelCase   =	qkv_bias
								UpperCamelCase   =	hidden_dropout_prob
								UpperCamelCase   =	attention_probs_dropout_prob
								UpperCamelCase   =	drop_path_rate
								UpperCamelCase   =	hidden_act
								UpperCamelCase   =	use_absolute_embeddings
								UpperCamelCase   =	layer_norm_eps
								UpperCamelCase   =	initializer_range
								UpperCamelCase   =	encoder_stride
								# we set the hidden_size attribute in order to make Swinv2 work with VisionEncoderDecoderModel
								# this indicates the channel dimension after the last stage of the model
								UpperCamelCase   =	int(embed_dim * 2 ** (len(__a		) - 1)		)
								UpperCamelCase   =	(0, 0, 0, 0)






