
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      _UpperCAmelCase				(					unittest.TestCase):
    def 					lowerCamelCase__	(    self     ):
      super().tearDown()
      gc.collect()
    def 					lowerCamelCase__	(    self     ):
      _snake_case     :						Any						=   FlaxStableDiffusionPipeline.from_pretrained(
          "stabilityai/stable-diffusion-2"						,    revision="bf16"						,    dtype=jnp.bfloataa						,    )

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

      _snake_case     :						Dict						=   replicate(_a     )
      _snake_case     :						Tuple						=   shard(_a     )

      _snake_case     :						int						=   jax.random.PRNGKey(0     )
      _snake_case     :						Optional[int]						=   jax.random.split(_a						,    jax.device_count()     )

      _snake_case     :						int						=   sd_pipe(_a						,    _a						,    _a						,    num_inference_steps=25						,    jit=_a     )[0]
      assert images.shape == (jax.device_count(), 1, 7_68, 7_68, 3)

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

      _snake_case     :						Optional[int]						=   jnp.asarray(jax.device_get(image_slice.flatten()     )     )
      _snake_case     :						Optional[int]						=   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 					lowerCamelCase__	(    self     ):
      _snake_case     :						str						=   """stabilityai/stable-diffusion-2"""
      _snake_case     :						Tuple						=   FlaxDPMSolverMultistepScheduler.from_pretrained(_a						,    subfolder="scheduler"     )
      _snake_case     :						int						=   FlaxStableDiffusionPipeline.from_pretrained(
          _a						,    scheduler=_a						,    revision="bf16"						,    dtype=jnp.bfloataa						,    )
      _snake_case     :						int						=   scheduler_params

      _snake_case     :						Tuple						=   """A painting of a squirrel eating a burger"""
      _snake_case     :						str						=   jax.device_count()
      _snake_case     :						int						=   num_samples * [prompt]
      _snake_case     :						Optional[int]						=   sd_pipe.prepare_inputs(_a     )

      _snake_case     :						str						=   replicate(_a     )
      _snake_case     :						Optional[int]						=   shard(_a     )

      _snake_case     :						List[Any]						=   jax.random.PRNGKey(0     )
      _snake_case     :						Optional[int]						=   jax.random.split(_a						,    jax.device_count()     )

      _snake_case     :						Tuple						=   sd_pipe(_a						,    _a						,    _a						,    num_inference_steps=25						,    jit=_a     )[0]
      assert images.shape == (jax.device_count(), 1, 7_68, 7_68, 3)

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

      _snake_case     :						int						=   jnp.asarray(jax.device_get(image_slice.flatten()     )     )
      _snake_case     :						Tuple						=   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
from ...utils.backbone_utils import BackboneConfigMixin, get_aligned_output_features_output_indices


_a      :			Optional[int]					=							logging.get_logger(__name__)

_a      :			List[str]					=							{
    """facebook/convnextv2-tiny-1k-224""": """https://huggingface.co/facebook/convnextv2-tiny-1k-224/resolve/main/config.json""",
}






class      _UpperCAmelCase				(					_snake_case		,					_snake_case):
    __lowercase							: List[Any]  =					"""convnextv2"""
    def __init__(    self						,    snake_case_=3						,    snake_case_=4						,    snake_case_=4						,    snake_case_=None						,    snake_case_=None						,    snake_case_="gelu"						,    snake_case_=0.02						,    snake_case_=1E-12						,    snake_case_=0.0						,    snake_case_=2_24						,    snake_case_=None						,    snake_case_=None						,    **snake_case_						,    ):
      super().__init__(**snake_case_     )

      _snake_case     :						Tuple						=   num_channels
      _snake_case     :						Optional[int]						=   patch_size
      _snake_case     :						Tuple						=   num_stages
      _snake_case     :						int						=   [96, 1_92, 3_84, 7_68] if hidden_sizes is None else hidden_sizes
      _snake_case     :						str						=   [3, 3, 9, 3] if depths is None else depths
      _snake_case     :						int						=   hidden_act
      _snake_case     :						Tuple						=   initializer_range
      _snake_case     :						Union[str, Any]						=   layer_norm_eps
      _snake_case     :						Optional[int]						=   drop_path_rate
      _snake_case     :						Union[str, Any]						=   image_size
      _snake_case     :						List[Any]						=   ["stem"] + [F'stage{idx}' for idx in range(1						,    len(self.depths     ) + 1     )]
      _snake_case			,					_snake_case     :						Dict						=   get_aligned_output_features_output_indices(
          out_features=snake_case_						,    out_indices=snake_case_						,    stage_names=self.stage_names     )

