

'''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


__snake_case	: List[str]       				=							logging.get_logger(__name__)

# General docstring
__snake_case	: List[Any]       				=							"""PoolFormerConfig"""

# Base docstring
__snake_case	: List[str]       				=							"""sail/poolformer_s12"""
__snake_case	: Optional[int]       				=							[1, 512, 7, 7]

# Image classification docstring
__snake_case	: Union[str, Any]       				=							"""sail/poolformer_s12"""
__snake_case	: List[str]       				=							"""tabby, tabby cat"""

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



def 					_lowercase						(   lowerCamelCase__			:		str,		lowerCamelCase__			:		float = 0.0,		lowerCamelCase__			:		bool = False							):



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





						def __init__(	self	,    snake_case_ = None			)			->							Optional[int]:
										super().__init__()
										_a        =   drop_prob





						def __lowerCAmelCase   (	self	,    snake_case_			)			->							Tuple:
										return drop_path(snake_case__	,    self.drop_prob	,    self.training			)





						def __lowerCAmelCase   (	self			)			->							str:
										return "p={}".format(self.drop_prob			)



class        A  (	nn.Module ):





						def __init__(	self	,    snake_case_	,    snake_case_	,    snake_case_	,    snake_case_	,    snake_case_	,    snake_case_=None			)			->							List[str]:
										super().__init__()
										_a        =   patch_size if isinstance(snake_case__	,    collections.abc.Iterable			) else (patch_size, patch_size)
										_a        =   stride if isinstance(snake_case__	,    collections.abc.Iterable			) else (stride, stride)
										_a        =   padding if isinstance(snake_case__	,    collections.abc.Iterable			) else (padding, padding)

										_a        =   nn.Convad(snake_case__	,    snake_case__	,    kernel_size=snake_case__	,    stride=snake_case__	,    padding=snake_case__			)
										_a        =   norm_layer(snake_case__			) if norm_layer else nn.Identity()





						def __lowerCAmelCase   (	self	,    snake_case_			)			->							Optional[Any]:
										_a        =   self.projection(snake_case__			)
										_a        =   self.norm(snake_case__			)
										return embeddings



class        A  (	nn.GroupNorm ):





						def __init__(	self	,    snake_case_	,    **snake_case_			)			->							Tuple:
										super().__init__(1	,    snake_case__	,    **snake_case__			)



class        A  (	nn.Module ):





						def __init__(	self	,    snake_case_			)			->							Tuple:
										super().__init__()
										_a        =   nn.AvgPoolad(snake_case__	,    stride=1	,    padding=pool_size // 2	,    count_include_pad=snake_case__			)





						def __lowerCAmelCase   (	self	,    snake_case_			)			->							str:
										return self.pool(snake_case__			) - hidden_states



class        A  (	nn.Module ):





						def __init__(	self	,    snake_case_	,    snake_case_	,    snake_case_	,    snake_case_			)			->							Tuple:
										super().__init__()
										_a        =   nn.Convad(snake_case__	,    snake_case__	,    1			)
										_a        =   nn.Convad(snake_case__	,    snake_case__	,    1			)
										_a        =   PoolFormerDropPath(snake_case__			)
										if isinstance(config.hidden_act	,    snake_case__			):
														_a        =   ACTaFN[config.hidden_act]
										else:
														_a        =   config.hidden_act





						def __lowerCAmelCase   (	self	,    snake_case_			)			->							List[str]:
										_a        =   self.conva(snake_case__			)
										_a        =   self.act_fn(snake_case__			)
										_a        =   self.drop(snake_case__			)
										_a        =   self.conva(snake_case__			)
										_a        =   self.drop(snake_case__			)

										return hidden_states



class        A  (	nn.Module ):





						def __init__(	self	,    snake_case_	,    snake_case_	,    snake_case_	,    snake_case_	,    snake_case_	,    snake_case_			)			->							List[Any]:
										super().__init__()
										_a        =   PoolFormerPooling(snake_case__			)
										_a        =   PoolFormerOutput(snake_case__	,    snake_case__	,    snake_case__	,    snake_case__			)
										_a        =   PoolFormerGroupNorm(snake_case__			)
										_a        =   PoolFormerGroupNorm(snake_case__			)

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





						def __lowerCAmelCase   (	self	,    snake_case_			)			->							List[str]:
										if self.use_layer_scale:
														_a        =   self.pooling(self.before_norm(snake_case__			)			)
														_a        =   self.layer_scale_a.unsqueeze(-1			).unsqueeze(-1			) * pooling_output
														# First residual connection
														_a        =   hidden_states + self.drop_path(snake_case__			)
														_a        =   ()

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

														_a        =   (output,) + outputs
														return outputs

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

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

														_a        =   (output,) + outputs
														return outputs



class        A  (	nn.Module ):





						def __init__(	self	,    snake_case_			)			->							Dict:
										super().__init__()
										_a        =   config
										# stochastic depth decay rule
										_a        =   [x.item() for x in torch.linspace(0	,    config.drop_path_rate	,    sum(config.depths			)			)]

										# patch embeddings
										_a        =   []
										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]	,    )			)
										_a        =   nn.ModuleList(snake_case__			)

										# Transformer blocks
										_a        =   []
										_a        =   0
										for i in range(config.num_encoder_blocks			):
														# each block consists of layers
														_a        =   []
														if i != 0:
																		cur += config.depths[i - 1]
														for j in range(config.depths[i]			):
																		layers.append(
																		    PoolFormerLayer(
																		        snake_case__	,    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(snake_case__			)			)

										_a        =   nn.ModuleList(snake_case__			)





						def __lowerCAmelCase   (	self	,    snake_case_	,    snake_case_=False	,    snake_case_=True			)			->							Optional[Any]:
										_a        =   () if output_hidden_states else None

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

														if output_hidden_states:
																		_a        =   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=snake_case__	,    hidden_states=snake_case__			)



class        A  (	a ):
						__UpperCAmelCase				:	Dict   =   PoolFormerConfig
						__UpperCAmelCase				:	Dict   =   """poolformer"""
						__UpperCAmelCase				:	Union[str, Any]   =   """pixel_values"""
						__UpperCAmelCase				:	str   =   True





						def __lowerCAmelCase   (	self	,    snake_case_			)			->							int:
										if isinstance(snake_case__	,    (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(snake_case__	,    nn.LayerNorm			):
														module.bias.data.zero_()
														module.weight.data.fill_(1.0			)





						def __lowerCAmelCase   (	self	,    snake_case_	,    snake_case_=False			)			->							Optional[Any]:
										if isinstance(snake_case__	,    snake_case__			):
														_a        =   value



__snake_case	: Optional[int]       				=							r"""
    This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) sub-class. Use
    it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and
    behavior.

    Parameters:
        config ([`PoolFormerConfig`]): Model configuration class with all the parameters of the model.
            Initializing with a config file does not load the weights associated with the model, only the
            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
"""

__snake_case	: Tuple       				=							r"""
    Args:
        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See
            [`PoolFormerImageProcessor.__call__`] for details.
"""
@add_start_docstrings(
    """The bare PoolFormer Model transformer outputting raw hidden-states without any specific head on top."""  ,    a  ,    )
class        A  (	a ):





						def __init__(	self	,    snake_case_			)			->							List[str]:
										super().__init__(snake_case__			)
										_a        =   config

										_a        =   PoolFormerEncoder(snake_case__			)

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





						def __lowerCAmelCase   (	self			)			->							Optional[int]:
										return self.embeddings.patch_embeddings





						@add_start_docstrings_to_model_forward(snake_case__			)
						@add_code_sample_docstrings(
						    checkpoint=_CHECKPOINT_FOR_DOC	,    output_type=snake_case__	,    config_class=_CONFIG_FOR_DOC	,    modality="vision"	,    expected_output=_EXPECTED_OUTPUT_SHAPE	,    )
						def __lowerCAmelCase   (	self	,    snake_case_ = None	,    snake_case_ = None	,    snake_case_ = None	,    )			->							Any:
										_a        =   (
										    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
										)
										_a        =   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"			)

										_a        =   self.encoder(
										    snake_case__	,    output_hidden_states=snake_case__	,    return_dict=snake_case__	,    )
										_a        =   encoder_outputs[0]

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

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



class        A  (	nn.Module ):





						def __init__(	self	,    snake_case_			)			->							Any:
										super().__init__()
										_a        =   nn.Linear(config.hidden_size	,    config.hidden_size			)





						def __lowerCAmelCase   (	self	,    snake_case_			)			->							str:
										_a        =   self.dense(snake_case__			)
										return output



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





						def __init__(	self	,    snake_case_			)			->							List[str]:
										super().__init__(snake_case__			)
										_a        =   config.num_labels
										_a        =   PoolFormerModel(snake_case__			)

										# Final norm
										_a        =   PoolFormerGroupNorm(config.hidden_sizes[-1]			)
										# Classifier head
										_a        =   (
										    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(snake_case__			)
						@add_code_sample_docstrings(
						    checkpoint=_IMAGE_CLASS_CHECKPOINT	,    output_type=snake_case__	,    config_class=_CONFIG_FOR_DOC	,    expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT	,    )
						def __lowerCAmelCase   (	self	,    snake_case_ = None	,    snake_case_ = None	,    snake_case_ = None	,    snake_case_ = None	,    )			->							Optional[Any]:
										_a        =   return_dict if return_dict is not None else self.config.use_return_dict

										_a        =   self.poolformer(
										    snake_case__	,    output_hidden_states=snake_case__	,    return_dict=snake_case__	,    )

										_a        =   outputs[0]

										_a        =   self.classifier(self.norm(snake_case__			).mean([-2, -1]			)			)

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

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

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

										return ImageClassifierOutputWithNoAttention(loss=snake_case__	,    logits=snake_case__	,    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 			lowerCAmelCase__			(    unittest.TestCase   ):
   def       __a							(		self      :	Optional[int]							):


        '''simple docstring'''
        # clean up the VRAM after each test
        super().tearDown()
        gc.collect()




   def       __a							(		self      :	List[str]							):


        '''simple docstring'''
        UpperCAmelCase__						:    Optional[int]							=							load_image(
            "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
            "/sd2-inpaint/init_image.png"							)
        UpperCAmelCase__						:    Any							=							load_image(
            "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-inpaint/mask.png"							)

        UpperCAmelCase__						:    str							=							"xvjiarui/stable-diffusion-2-inpainting"
        UpperCAmelCase__							,				UpperCAmelCase__						:    List[Any]							=							FlaxStableDiffusionInpaintPipeline.from_pretrained(snake_case__							,			safety_checker=snake_case__							)

        UpperCAmelCase__						:    List[str]							=							"Face of a yellow cat, high resolution, sitting on a park bench"

        UpperCAmelCase__						:    List[str]							=							jax.random.PRNGKey(0							)
        UpperCAmelCase__						:    Optional[int]							=							5_0

        UpperCAmelCase__						:    List[Any]							=							jax.device_count()
        UpperCAmelCase__						:    int							=							num_samples * [prompt]
        UpperCAmelCase__						:    str							=							num_samples * [init_image]
        UpperCAmelCase__						:    Optional[Any]							=							num_samples * [mask_image]
        UpperCAmelCase__							,				UpperCAmelCase__							,				UpperCAmelCase__						:    Optional[int]							=							pipeline.prepare_inputs(snake_case__							,			snake_case__							,			snake_case__							)

        # shard inputs and rng
        UpperCAmelCase__						:    Dict							=							replicate(snake_case__							)
        UpperCAmelCase__						:    Dict							=							jax.random.split(snake_case__							,			jax.device_count()							)
        UpperCAmelCase__						:    List[str]							=							shard(snake_case__							)
        UpperCAmelCase__						:    str							=							shard(snake_case__							)
        UpperCAmelCase__						:    Optional[int]							=							shard(snake_case__							)

        UpperCAmelCase__						:    Dict							=							pipeline(
            snake_case__							,			snake_case__							,			snake_case__							,			snake_case__							,			snake_case__							,			snake_case__							,			jit=snake_case__							)

        UpperCAmelCase__						:    List[str]							=							output.images.reshape(snake_case__							,			5_1_2							,			5_1_2							,			3							)

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

        UpperCAmelCase__						:    str							=							jnp.asarray(jax.device_get(image_slice.flatten()							)							)
        UpperCAmelCase__						:    int							=							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


