



'''simple docstring'''



import collections.abc
from typing import Optional, Tuple, Union

import torch
import torch.utils.checkpoint
from torch import nn
from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss

from ...activations import ACTaFN
from ...modeling_outputs import BaseModelOutputWithNoAttention, ImageClassifierOutputWithNoAttention
from ...modeling_utils import PreTrainedModel
from ...utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward, logging
from .configuration_poolformer import PoolFormerConfig


lowerCamelCase	 =       logging.get_logger(__name__)

# General docstring
lowerCamelCase	 =       """PoolFormerConfig"""

# Base docstring
lowerCamelCase	 =       """sail/poolformer_s12"""
lowerCamelCase	 =       [1, 512, 7, 7]

# Image classification docstring
lowerCamelCase	 =       """sail/poolformer_s12"""
lowerCamelCase	 =       """tabby, tabby cat"""

lowerCamelCase	 =       [
    """sail/poolformer_s12""",
    # See all PoolFormer models at https://huggingface.co/models?filter=poolformer
]



def   _A       (   _lowerCAmelCase  ,   _lowerCAmelCase = 0.0  ,   _lowerCAmelCase = False					):



							"""simple docstring"""



							if drop_prob == 0.0 or not training:
														return input
							__lowercase       	=1 - drop_prob
							__lowercase       	=(input.shape[0],) + (1,) * (input.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
							__lowercase       	=keep_prob + torch.rand(lowerCamelCase_  ,   dtype=input.dtype  ,   device=input.device					)
							random_tensor.floor_()  # binarize
							__lowercase       	=input.div(lowerCamelCase_					) * random_tensor
							return output





class _UpperCamelCase		(     nn.Module		):





				'''simple docstring'''





				def __init__(							self  :  Optional[Any]  ,       _lowerCAmelCase  :  Optional[float] = None):







											'''simple docstring'''

											super().__init__()
											__lowercase       	=drop_prob

				def 				__lowerCamelCase     (							self  :  List[str]  ,       _lowerCAmelCase  :  torch.Tensor):







											'''simple docstring'''

											return drop_path(UpperCAmelCase_  ,       self.drop_prob  ,       self.training)

				def 				__lowerCamelCase     (							self  :  Any):







											'''simple docstring'''

											return "p={}".format(self.drop_prob)





class _UpperCamelCase		(     nn.Module		):





				'''simple docstring'''





				def __init__(							self  :  Union[str, Any]  ,       _lowerCAmelCase  :  Union[str, Any]  ,       _lowerCAmelCase  :  str  ,       _lowerCAmelCase  :  Tuple  ,       _lowerCAmelCase  :  Tuple  ,       _lowerCAmelCase  :  Union[str, Any]  ,       _lowerCAmelCase  :  Optional[Any]=None):







											'''simple docstring'''

											super().__init__()
											__lowercase       	=patch_size if isinstance(UpperCAmelCase_  ,       collections.abc.Iterable) else (patch_size, patch_size)
											__lowercase       	=stride if isinstance(UpperCAmelCase_  ,       collections.abc.Iterable) else (stride, stride)
											__lowercase       	=padding if isinstance(UpperCAmelCase_  ,       collections.abc.Iterable) else (padding, padding)

											__lowercase       	=nn.Convad(UpperCAmelCase_  ,       UpperCAmelCase_  ,       kernel_size=UpperCAmelCase_  ,       stride=UpperCAmelCase_  ,       padding=UpperCAmelCase_)
											__lowercase       	=norm_layer(UpperCAmelCase_) if norm_layer else nn.Identity()

				def 				__lowerCamelCase     (							self  :  str  ,       _lowerCAmelCase  :  List[str]):







											'''simple docstring'''

											__lowercase       	=self.projection(UpperCAmelCase_)
											__lowercase       	=self.norm(UpperCAmelCase_)
											return embeddings





class _UpperCamelCase		(     nn.GroupNorm		):





				'''simple docstring'''





				def __init__(							self  :  Optional[int]  ,       _lowerCAmelCase  :  Optional[Any]  ,       **_lowerCAmelCase  :  Tuple):







											'''simple docstring'''

											super().__init__(1  ,       UpperCAmelCase_  ,       **UpperCAmelCase_)





class _UpperCamelCase		(     nn.Module		):





				'''simple docstring'''





				def __init__(							self  :  Union[str, Any]  ,       _lowerCAmelCase  :  Any):







											'''simple docstring'''

											super().__init__()
											__lowercase       	=nn.AvgPoolad(UpperCAmelCase_  ,       stride=1  ,       padding=pool_size // 2  ,       count_include_pad=UpperCAmelCase_)

				def 				__lowerCamelCase     (							self  :  Any  ,       _lowerCAmelCase  :  Optional[int]):







											'''simple docstring'''

											return self.pool(UpperCAmelCase_) - hidden_states





class _UpperCamelCase		(     nn.Module		):





				'''simple docstring'''





				def __init__(							self  :  int  ,       _lowerCAmelCase  :  int  ,       _lowerCAmelCase  :  Any  ,       _lowerCAmelCase  :  List[Any]  ,       _lowerCAmelCase  :  List[Any]):







											'''simple docstring'''

											super().__init__()
											__lowercase       	=nn.Convad(UpperCAmelCase_  ,       UpperCAmelCase_  ,       1)
											__lowercase       	=nn.Convad(UpperCAmelCase_  ,       UpperCAmelCase_  ,       1)
											__lowercase       	=PoolFormerDropPath(UpperCAmelCase_)
											if isinstance(config.hidden_act  ,       UpperCAmelCase_):
																		__lowercase       	=ACTaFN[config.hidden_act]
											else:
																		__lowercase       	=config.hidden_act

				def 				__lowerCamelCase     (							self  :  Optional[int]  ,       _lowerCAmelCase  :  Tuple):







											'''simple docstring'''

											__lowercase       	=self.conva(UpperCAmelCase_)
											__lowercase       	=self.act_fn(UpperCAmelCase_)
											__lowercase       	=self.drop(UpperCAmelCase_)
											__lowercase       	=self.conva(UpperCAmelCase_)
											__lowercase       	=self.drop(UpperCAmelCase_)

											return hidden_states





class _UpperCamelCase		(     nn.Module		):





				'''simple docstring'''





				def __init__(							self  :  List[Any]  ,       _lowerCAmelCase  :  List[str]  ,       _lowerCAmelCase  :  Union[str, Any]  ,       _lowerCAmelCase  :  int  ,       _lowerCAmelCase  :  List[Any]  ,       _lowerCAmelCase  :  Dict  ,       _lowerCAmelCase  :  Any):







											'''simple docstring'''

											super().__init__()
											__lowercase       	=PoolFormerPooling(UpperCAmelCase_)
											__lowercase       	=PoolFormerOutput(UpperCAmelCase_  ,       UpperCAmelCase_  ,       UpperCAmelCase_  ,       UpperCAmelCase_)
											__lowercase       	=PoolFormerGroupNorm(UpperCAmelCase_)
											__lowercase       	=PoolFormerGroupNorm(UpperCAmelCase_)

											# Useful for training neural nets
											__lowercase       	=PoolFormerDropPath(UpperCAmelCase_) if drop_path > 0.0 else nn.Identity()
											__lowercase       	=config.use_layer_scale
											if config.use_layer_scale:
																		__lowercase       	=nn.Parameter(
																		    config.layer_scale_init_value * torch.ones((UpperCAmelCase_))  ,       requires_grad=UpperCAmelCase_)
																		__lowercase       	=nn.Parameter(
																		    config.layer_scale_init_value * torch.ones((UpperCAmelCase_))  ,       requires_grad=UpperCAmelCase_)

				def 				__lowerCamelCase     (							self  :  Tuple  ,       _lowerCAmelCase  :  str):







											'''simple docstring'''

											if self.use_layer_scale:
																		__lowercase       	=self.pooling(self.before_norm(UpperCAmelCase_))
																		__lowercase       	=self.layer_scale_a.unsqueeze(-1).unsqueeze(-1) * pooling_output
																		# First residual connection
																		__lowercase       	=hidden_states + self.drop_path(UpperCAmelCase_)
																		__lowercase       	=()

																		__lowercase       	=self.output(self.after_norm(UpperCAmelCase_))
																		__lowercase       	=self.layer_scale_a.unsqueeze(-1).unsqueeze(-1) * layer_output
																		# Second residual connection
																		__lowercase       	=hidden_states + self.drop_path(UpperCAmelCase_)

																		__lowercase       	=(output,) + outputs
																		return outputs

											else:
																		__lowercase       	=self.drop_path(self.pooling(self.before_norm(UpperCAmelCase_)))
																		# First residual connection
																		__lowercase       	=pooling_output + hidden_states
																		__lowercase       	=()

																		# Second residual connection inside the PoolFormerOutput block
																		__lowercase       	=self.drop_path(self.output(self.after_norm(UpperCAmelCase_)))
																		__lowercase       	=hidden_states + layer_output

																		__lowercase       	=(output,) + outputs
																		return outputs





class _UpperCamelCase		(     nn.Module		):





				'''simple docstring'''





				def __init__(							self  :  Dict  ,       _lowerCAmelCase  :  Optional[Any]):







											'''simple docstring'''

											super().__init__()
											__lowercase       	=config
											# stochastic depth decay rule
											__lowercase       	=[x.item() for x in torch.linspace(0  ,       config.drop_path_rate  ,       sum(config.depths))]

											# patch embeddings
											__lowercase       	=[]
											for i in range(config.num_encoder_blocks):
																		embeddings.append(
																		    PoolFormerEmbeddings(
																		        patch_size=config.patch_sizes[i]  ,       stride=config.strides[i]  ,       padding=config.padding[i]  ,       num_channels=config.num_channels if i == 0 else config.hidden_sizes[i - 1]  ,       hidden_size=config.hidden_sizes[i]  ,       ))
											__lowercase       	=nn.ModuleList(UpperCAmelCase_)

											# Transformer blocks
											__lowercase       	=[]
											__lowercase       	=0
											for i in range(config.num_encoder_blocks):
																		# each block consists of layers
																		__lowercase       	=[]
																		if i != 0:
																									cur += config.depths[i - 1]
																		for j in range(config.depths[i]):
																									layers.append(
																									    PoolFormerLayer(
																									        UpperCAmelCase_  ,       num_channels=config.hidden_sizes[i]  ,       pool_size=config.pool_size  ,       hidden_size=config.hidden_sizes[i]  ,       intermediate_size=int(config.hidden_sizes[i] * config.mlp_ratio)  ,       drop_path=dpr[cur + j]  ,       ))
																		blocks.append(nn.ModuleList(UpperCAmelCase_))

											__lowercase       	=nn.ModuleList(UpperCAmelCase_)

				def 				__lowerCamelCase     (							self  :  Dict  ,       _lowerCAmelCase  :  Any  ,       _lowerCAmelCase  :  Optional[Any]=False  ,       _lowerCAmelCase  :  int=True):







											'''simple docstring'''

											__lowercase       	=() if output_hidden_states else None

											__lowercase       	=pixel_values
											for idx, layers in enumerate(zip(self.patch_embeddings  ,       self.block)):
																		__lowercase       	=layers
																		# Get patch embeddings from hidden_states
																		__lowercase       	=embedding_layer(UpperCAmelCase_)
																		# Send the embeddings through the blocks
																		for _, blk in enumerate(UpperCAmelCase_):
																									__lowercase       	=blk(UpperCAmelCase_)
																									__lowercase       	=layer_outputs[0]

																		if output_hidden_states:
																									__lowercase       	=all_hidden_states + (hidden_states,)

											if not return_dict:
																		return tuple(v for v in [hidden_states, all_hidden_states] if v is not None)

											return BaseModelOutputWithNoAttention(last_hidden_state=UpperCAmelCase_  ,       hidden_states=UpperCAmelCase_)





class _UpperCamelCase		(     __lowerCamelCase		):





				'''simple docstring'''





				lowerCAmelCase__					    =    PoolFormerConfig
				lowerCAmelCase__					    =    """poolformer"""
				lowerCAmelCase__					    =    """pixel_values"""
				lowerCAmelCase__					    =    True

				def 				__lowerCamelCase     (							self  :  int  ,       _lowerCAmelCase  :  Dict):







											'''simple docstring'''

											if isinstance(UpperCAmelCase_  ,       (nn.Linear, nn.Convad)):
																		module.weight.data.normal_(mean=0.0  ,       std=self.config.initializer_range)
																		if module.bias is not None:
																									module.bias.data.zero_()
											elif isinstance(UpperCAmelCase_  ,       nn.LayerNorm):
																		module.bias.data.zero_()
																		module.weight.data.fill_(1.0)

				def 				__lowerCamelCase     (							self  :  int  ,       _lowerCAmelCase  :  Optional[Any]  ,       _lowerCAmelCase  :  Optional[Any]=False):







											'''simple docstring'''

											if isinstance(UpperCAmelCase_  ,       UpperCAmelCase_):
																		__lowercase       	=value

lowerCamelCase	 =       R"""\n    This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) sub-class. Use\n    it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and\n    behavior.\n\n    Parameters:\n        config ([`PoolFormerConfig`]): Model configuration class with all the parameters of the model.\n            Initializing with a config file does not load the weights associated with the model, only the\n            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.\n"""

lowerCamelCase	 =       R"""\n    Args:\n        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):\n            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See\n            [`PoolFormerImageProcessor.__call__`] for details.\n"""





@add_start_docstrings(
    """The bare PoolFormer Model transformer outputting raw hidden-states without any specific head on top."""      ,       __lowerCamelCase      ,       )
class _UpperCamelCase		(     __lowerCamelCase		):





				'''simple docstring'''





				def __init__(							self  :  Tuple  ,       _lowerCAmelCase  :  int):







											'''simple docstring'''

											super().__init__(UpperCAmelCase_)
											__lowercase       	=config

											__lowercase       	=PoolFormerEncoder(UpperCAmelCase_)

											# Initialize weights and apply final processing
											self.post_init()

				def 				__lowerCamelCase     (							self  :  int):







											'''simple docstring'''

											return self.embeddings.patch_embeddings

				@add_start_docstrings_to_model_forward(UpperCAmelCase_)
				@add_code_sample_docstrings(
				    checkpoint=_CHECKPOINT_FOR_DOC  ,       output_type=UpperCAmelCase_  ,       config_class=_CONFIG_FOR_DOC  ,       modality='vision'  ,       expected_output=_EXPECTED_OUTPUT_SHAPE  ,       )
				def 				__lowerCamelCase     (							self  :  Union[str, Any]  ,       _lowerCAmelCase  :  Optional[torch.FloatTensor] = None  ,       _lowerCAmelCase  :  Optional[bool] = None  ,       _lowerCAmelCase  :  Optional[bool] = None  ,       ):







											'''simple docstring'''

											__lowercase       	=(
											    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
											)
											__lowercase       	=return_dict if return_dict is not None else self.config.use_return_dict

											if pixel_values is None:
																		raise ValueError('You have to specify pixel_values')

											__lowercase       	=self.encoder(
											    UpperCAmelCase_  ,       output_hidden_states=UpperCAmelCase_  ,       return_dict=UpperCAmelCase_  ,       )
											__lowercase       	=encoder_outputs[0]

											if not return_dict:
																		return (sequence_output, None) + encoder_outputs[1:]

											return BaseModelOutputWithNoAttention(
											    last_hidden_state=UpperCAmelCase_  ,       hidden_states=encoder_outputs.hidden_states  ,       )





class _UpperCamelCase		(     nn.Module		):





				'''simple docstring'''





				def __init__(							self  :  Union[str, Any]  ,       _lowerCAmelCase  :  Tuple):







											'''simple docstring'''

											super().__init__()
											__lowercase       	=nn.Linear(config.hidden_size  ,       config.hidden_size)

				def 				__lowerCamelCase     (							self  :  List[Any]  ,       _lowerCAmelCase  :  List[Any]):







											'''simple docstring'''

											__lowercase       	=self.dense(UpperCAmelCase_)
											return output





@add_start_docstrings(
    """
    PoolFormer Model transformer with an image classification head on top
    """      ,       __lowerCamelCase      ,       )
class _UpperCamelCase		(     __lowerCamelCase		):





				'''simple docstring'''





				def __init__(							self  :  Optional[int]  ,       _lowerCAmelCase  :  Dict):







											'''simple docstring'''

											super().__init__(UpperCAmelCase_)
											__lowercase       	=config.num_labels
											__lowercase       	=PoolFormerModel(UpperCAmelCase_)

											# Final norm
											__lowercase       	=PoolFormerGroupNorm(config.hidden_sizes[-1])
											# Classifier head
											__lowercase       	=(
											    nn.Linear(config.hidden_sizes[-1]  ,       config.num_labels) if config.num_labels > 0 else nn.Identity()
											)

											# Initialize weights and apply final processing
											self.post_init()

				@add_start_docstrings_to_model_forward(UpperCAmelCase_)
				@add_code_sample_docstrings(
				    checkpoint=_IMAGE_CLASS_CHECKPOINT  ,       output_type=UpperCAmelCase_  ,       config_class=_CONFIG_FOR_DOC  ,       expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT  ,       )
				def 				__lowerCamelCase     (							self  :  List[str]  ,       _lowerCAmelCase  :  Optional[torch.FloatTensor] = None  ,       _lowerCAmelCase  :  Optional[torch.LongTensor] = None  ,       _lowerCAmelCase  :  Optional[bool] = None  ,       _lowerCAmelCase  :  Optional[bool] = None  ,       ):







											'''simple docstring'''

											__lowercase       	=return_dict if return_dict is not None else self.config.use_return_dict

											__lowercase       	=self.poolformer(
											    UpperCAmelCase_  ,       output_hidden_states=UpperCAmelCase_  ,       return_dict=UpperCAmelCase_  ,       )

											__lowercase       	=outputs[0]

											__lowercase       	=self.classifier(self.norm(UpperCAmelCase_).mean([-2, -1]))

											__lowercase       	=None
											if labels is not None:
																		if self.config.problem_type is None:
																									if self.num_labels == 1:
																																__lowercase       	='regression'
																									elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
																																__lowercase       	='single_label_classification'
																									else:
																																__lowercase       	='multi_label_classification'

																		if self.config.problem_type == "regression":
																									__lowercase       	=MSELoss()
																									if self.num_labels == 1:
																																__lowercase       	=loss_fct(logits.squeeze()  ,       labels.squeeze())
																									else:
																																__lowercase       	=loss_fct(UpperCAmelCase_  ,       UpperCAmelCase_)
																		elif self.config.problem_type == "single_label_classification":
																									__lowercase       	=CrossEntropyLoss()
																									__lowercase       	=loss_fct(logits.view(-1  ,       self.num_labels)  ,       labels.view(-1))
																		elif self.config.problem_type == "multi_label_classification":
																									__lowercase       	=BCEWithLogitsLoss()
																									__lowercase       	=loss_fct(UpperCAmelCase_  ,       UpperCAmelCase_)

											if not return_dict:
																		__lowercase       	=(logits,) + outputs[2:]
																		return ((loss,) + output) if loss is not None else output

											return ImageClassifierOutputWithNoAttention(loss=UpperCAmelCase_  ,       logits=UpperCAmelCase_  ,       hidden_states=outputs.hidden_states)








'''simple docstring'''



import gc
import unittest

from diffusers import FlaxStableDiffusionInpaintPipeline
from diffusers.utils import is_flax_available, load_image, 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		):





				'''simple docstring'''





				def 				__lowerCamelCase     (							self  :  List[Any]):







											'''simple docstring'''

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

				def 				__lowerCamelCase     (							self  :  List[str]):







											'''simple docstring'''

											__lowercase       	=load_image(
											    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
											    '/sd2-inpaint/init_image.png')
											__lowercase       	=load_image(
											    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-inpaint/mask.png')

											__lowercase       	='xvjiarui/stable-diffusion-2-inpainting'
											__lowercase	,				__lowercase       	=FlaxStableDiffusionInpaintPipeline.from_pretrained(_lowerCAmelCase  ,       safety_checker=_lowerCAmelCase)

											__lowercase       	='Face of a yellow cat, high resolution, sitting on a park bench'

											__lowercase       	=jax.random.PRNGKey(0)
											__lowercase       	=5_0

											__lowercase       	=jax.device_count()
											__lowercase       	=num_samples * [prompt]
											__lowercase       	=num_samples * [init_image]
											__lowercase       	=num_samples * [mask_image]
											__lowercase	,				__lowercase	,				__lowercase       	=pipeline.prepare_inputs(_lowerCAmelCase  ,       _lowerCAmelCase  ,       _lowerCAmelCase)

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

											__lowercase       	=pipeline(
											    _lowerCAmelCase  ,       _lowerCAmelCase  ,       _lowerCAmelCase  ,       _lowerCAmelCase  ,       _lowerCAmelCase  ,       _lowerCAmelCase  ,       jit=_lowerCAmelCase)

											__lowercase       	=output.images.reshape(_lowerCAmelCase  ,       5_1_2  ,       5_1_2  ,       3)

											__lowercase       	=images[0, 2_5_3:2_5_6, 2_5_3:2_5_6, -1]

											__lowercase       	=jnp.asarray(jax.device_get(image_slice.flatten()))
											__lowercase       	=jnp.array(
											    [0.361_1307, 0.3764_9736, 0.375_7408, 0.3821_3953, 0.3929_5167, 0.384_1631, 0.4155_4978, 0.413_7475, 0.421_7084])
											print(f"""output_slice: {output_slice}""")
											assert jnp.abs(output_slice - expected_slice).max() < 1e-2



