



import gc
import random
import unittest

import numpy as np
import torch

from diffusers import DDIMScheduler, KandinskyVaaPipeline, KandinskyVaaPriorPipeline, UNetaDConditionModel, VQModel
from diffusers.utils import floats_tensor, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()


class 				_a						(       SCREAMING_SNAKE_CASE       ,  unittest.TestCase			):





							'''simple docstring'''

							A				:Tuple =       KandinskyVaaPipeline
							A				:str =       [
							    "image_embeds",
							    "negative_image_embeds",
							]
							A				:Optional[Any] =       ["image_embeds", "negative_image_embeds"]
							A				:List[str] =       [
							    "generator",
							    "height",
							    "width",
							    "latents",
							    "guidance_scale",
							    "num_inference_steps",
							    "return_dict",
							    "guidance_scale",
							    "num_images_per_prompt",
							    "output_type",
							    "return_dict",
							]
							A				:List[str] =       False



							@property
							def 		_A	(							self				):







											"""simple docstring"""



											return 32



							@property
							def 		_A	(							self				):







											"""simple docstring"""



											return 32



							@property
							def 		_A	(							self				):







											"""simple docstring"""



											return self.time_input_dim



							@property
							def 		_A	(							self				):







											"""simple docstring"""



											return self.time_input_dim * 4



							@property
							def 		_A	(							self				):







											"""simple docstring"""



											return 100



							@property
							def 		_A	(							self				):







											"""simple docstring"""



											torch.manual_seed(0				)

											a__	:							List[Any]      						= {
											    "in_channels": 4,
											    # Out channels is double in channels because predicts mean and variance
											    "out_channels": 8,
											    "addition_embed_type": "image",
											    "down_block_types": ("ResnetDownsampleBlock2D", "SimpleCrossAttnDownBlock2D"),
											    "up_block_types": ("SimpleCrossAttnUpBlock2D", "ResnetUpsampleBlock2D"),
											    "mid_block_type": "UNetMidBlock2DSimpleCrossAttn",
											    "block_out_channels": (self.block_out_channels_a, self.block_out_channels_a * 2),
											    "layers_per_block": 1,
											    "encoder_hid_dim": self.text_embedder_hidden_size,
											    "encoder_hid_dim_type": "image_proj",
											    "cross_attention_dim": self.cross_attention_dim,
											    "attention_head_dim": 4,
											    "resnet_time_scale_shift": "scale_shift",
											    "class_embed_type": None,
											}

											a__	:							List[Any]      						= UNetaDConditionModel(**__UpperCAmelCase				)
											return model



							@property
							def 		_A	(							self				):







											"""simple docstring"""



											return {
											    "block_out_channels": [32, 64],
											    "down_block_types": ["DownEncoderBlock2D", "AttnDownEncoderBlock2D"],
											    "in_channels": 3,
											    "latent_channels": 4,
											    "layers_per_block": 1,
											    "norm_num_groups": 8,
											    "norm_type": "spatial",
											    "num_vq_embeddings": 12,
											    "out_channels": 3,
											    "up_block_types": [
											        "AttnUpDecoderBlock2D",
											        "UpDecoderBlock2D",
											    ],
											    "vq_embed_dim": 4,
											}



							@property
							def 		_A	(							self				):







											"""simple docstring"""



											torch.manual_seed(0				)
											a__	:							str      						= VQModel(**self.dummy_movq_kwargs				)
											return model



							def 		_A	(							self				):







											"""simple docstring"""



											a__	:							Dict      						= self.dummy_unet
											a__	:							Union[str, Any]      						= self.dummy_movq

											a__	:							int      						= DDIMScheduler(
											    num_train_timesteps=1000 ,					beta_schedule="linear" ,					beta_start=0.0_0_0_8_5 ,					beta_end=0.0_1_2 ,					clip_sample=__UpperCAmelCase ,					set_alpha_to_one=__UpperCAmelCase ,					steps_offset=1 ,					prediction_type="epsilon" ,					thresholding=__UpperCAmelCase ,					)

											a__	:							List[Any]      						= {
											    "unet": unet,
											    "scheduler": scheduler,
											    "movq": movq,
											}
											return components



							def 		_A	(							self ,					__UpperCAmelCase ,					__UpperCAmelCase=0				):







											"""simple docstring"""



											a__	:							List[str]      						= floats_tensor((1, self.text_embedder_hidden_size) ,					rng=random.Random(__UpperCAmelCase				)				).to(__UpperCAmelCase				)
											a__	:							Dict      						= floats_tensor((1, self.text_embedder_hidden_size) ,					rng=random.Random(seed + 1				)				).to(
											    __UpperCAmelCase				)
											if str(__UpperCAmelCase				).startswith("mps"				):
															a__	:							List[Any]      						= torch.manual_seed(__UpperCAmelCase				)
											else:
															a__	:							str      						= torch.Generator(device=__UpperCAmelCase				).manual_seed(__UpperCAmelCase				)
											a__	:							Dict      						= {
											    "image_embeds": image_embeds,
											    "negative_image_embeds": negative_image_embeds,
											    "generator": generator,
											    "height": 64,
											    "width": 64,
											    "guidance_scale": 4.0,
											    "num_inference_steps": 2,
											    "output_type": "np",
											}
											return inputs




							def 		_A	(							self				):







											"""simple docstring"""



											a__	:							int      						= "cpu"

											a__	:							List[Any]      						= self.get_dummy_components()

											a__	:							Union[str, Any]      						= self.pipeline_class(**__UpperCAmelCase				)
											a__	:							Any      						= pipe.to(__UpperCAmelCase				)

											pipe.set_progress_bar_config(disable=__UpperCAmelCase				)

											a__	:							List[Any]      						= pipe(**self.get_dummy_inputs(__UpperCAmelCase				)				)
											a__	:							Dict      						= output.images

											a__	:							Union[str, Any]      						= pipe(
											    **self.get_dummy_inputs(__UpperCAmelCase				) ,					return_dict=__UpperCAmelCase ,					)[0]

											a__	:							int      						= image[0, -3:, -3:, -1]
											a__	:							Optional[Any]      						= image_from_tuple[0, -3:, -3:, -1]

											assert image.shape == (1, 64, 64, 3)

											a__	:							List[Any]      						= np.array(
											    [0.6_2_3_7_9_7_6, 1.0, 0.3_6_4_4_1_3_3_2, 1.0, 0.7_0_6_3_9_6_3_4, 0.2_9_8_7_7_1_8_6, 0.8_5_6_5_2_1_2_5, 0.5_2_1_6_8_4_3, 0.5_4_4_5_4_0_4_6]				)

											assert (
											    np.abs(image_slice.flatten() - expected_slice				).max() < 1E-2
											), f' expected_slice {expected_slice}, but got {image_slice.flatten()}'

											assert (
											    np.abs(image_from_tuple_slice.flatten() - expected_slice				).max() < 1E-2
											), f' expected_slice {expected_slice}, but got {image_from_tuple_slice.flatten()}'





@slow
@require_torch_gpu
class 				_a						(       unittest.TestCase			):





							'''simple docstring'''

							def 		_A	(							self				):







											"""simple docstring"""



											super().tearDown()
											gc.collect()
											torch.cuda.empty_cache()




							def 		_A	(							self				):







											"""simple docstring"""



											a__	:							Optional[int]      						= load_numpy(
											    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
											    "/kandinskyv22/kandinskyv22_text2img_cat_fp16.npy"				)

											a__	:							str      						= KandinskyVaaPriorPipeline.from_pretrained(
											    "kandinsky-community/kandinsky-2-2-prior" ,					torch_dtype=torch.floataa				)
											pipe_prior.to(__UpperCAmelCase				)

											a__	:							Union[str, Any]      						= KandinskyVaaPipeline.from_pretrained(
											    "kandinsky-community/kandinsky-2-2-decoder" ,					torch_dtype=torch.floataa				)
											a__	:							Tuple      						= pipeline.to(__UpperCAmelCase				)
											pipeline.set_progress_bar_config(disable=__UpperCAmelCase				)

											a__	:							List[str]      						= "red cat, 4k photo"

											a__	:							str      						= torch.Generator(device="cuda"				).manual_seed(0				)
											a__						,   a__	:							Dict      						= pipe_prior(
											    __UpperCAmelCase ,					generator=__UpperCAmelCase ,					num_inference_steps=5 ,					negative_prompt="" ,					).to_tuple()

											a__	:							List[str]      						= torch.Generator(device="cuda"				).manual_seed(0				)
											a__	:							Union[str, Any]      						= pipeline(
											    image_embeds=__UpperCAmelCase ,					negative_image_embeds=__UpperCAmelCase ,					generator=__UpperCAmelCase ,					num_inference_steps=100 ,					output_type="np" ,					)

											a__	:							int      						= output.images[0]

											assert image.shape == (512, 512, 3)

											assert_mean_pixel_difference(__UpperCAmelCase ,					__UpperCAmelCase				)







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, 5_12, 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     SCREAMING_SNAKE_CASE(      __UpperCamelCase	,				__UpperCamelCase = 0.0	,				__UpperCamelCase = False )	->						Dict:
				if drop_prob == 0.0 or not training:
								return input
				a__	:							Tuple      						= 1 - drop_prob
				a__	:							Any      						= (input.shape[0],) + (1,) * (input.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
				a__	:							Dict      						= keep_prob + torch.rand(__UpperCamelCase	,				dtype=input.dtype	,				device=input.device )
				random_tensor.floor_()  # binarize
				a__	:							Optional[int]      						= input.div(__UpperCamelCase ) * random_tensor
				return output


class 				_a						(       nn.Module			):





							'''simple docstring'''

							def __init__(							self ,					__UpperCAmelCase = None				):







											"""simple docstring"""



											super().__init__()
											a__	:							Optional[Any]      						= drop_prob



							def 		_A	(							self ,					__UpperCAmelCase				):







											"""simple docstring"""



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




							def 		_A	(							self				):







											"""simple docstring"""



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





class 				_a						(       nn.Module			):





							'''simple docstring'''

							def __init__(							self ,					__UpperCAmelCase ,					__UpperCAmelCase ,					__UpperCAmelCase ,					__UpperCAmelCase ,					__UpperCAmelCase ,					__UpperCAmelCase=None				):







											"""simple docstring"""



											super().__init__()
											a__	:							Optional[Any]      						= patch_size if isinstance(__UpperCAmelCase ,					collections.abc.Iterable				) else (patch_size, patch_size)
											a__	:							List[str]      						= stride if isinstance(__UpperCAmelCase ,					collections.abc.Iterable				) else (stride, stride)
											a__	:							Union[str, Any]      						= padding if isinstance(__UpperCAmelCase ,					collections.abc.Iterable				) else (padding, padding)

											a__	:							int      						= nn.Convad(__UpperCAmelCase ,					__UpperCAmelCase ,					kernel_size=__UpperCAmelCase ,					stride=__UpperCAmelCase ,					padding=__UpperCAmelCase				)
											a__	:							Optional[Any]      						= norm_layer(__UpperCAmelCase				) if norm_layer else nn.Identity()




							def 		_A	(							self ,					__UpperCAmelCase				):







											"""simple docstring"""



											a__	:							Dict      						= self.projection(__UpperCAmelCase				)
											a__	:							Union[str, Any]      						= self.norm(__UpperCAmelCase				)
											return embeddings





class 				_a						(       nn.GroupNorm			):





							'''simple docstring'''

							def __init__(							self ,					__UpperCAmelCase ,					**__UpperCAmelCase				):







											"""simple docstring"""



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





class 				_a						(       nn.Module			):





							'''simple docstring'''

							def __init__(							self ,					__UpperCAmelCase				):







											"""simple docstring"""



											super().__init__()
											a__	:							List[str]      						= nn.AvgPoolad(__UpperCAmelCase ,					stride=1 ,					padding=pool_size // 2 ,					count_include_pad=__UpperCAmelCase				)




							def 		_A	(							self ,					__UpperCAmelCase				):







											"""simple docstring"""



											return self.pool(__UpperCAmelCase				) - hidden_states





class 				_a						(       nn.Module			):





							'''simple docstring'''

							def __init__(							self ,					__UpperCAmelCase ,					__UpperCAmelCase ,					__UpperCAmelCase ,					__UpperCAmelCase				):







											"""simple docstring"""



											super().__init__()
											a__	:							str      						= nn.Convad(__UpperCAmelCase ,					__UpperCAmelCase ,					1				)
											a__	:							int      						= nn.Convad(__UpperCAmelCase ,					__UpperCAmelCase ,					1				)
											a__	:							Any      						= PoolFormerDropPath(__UpperCAmelCase				)
											if isinstance(config.hidden_act ,					__UpperCAmelCase				):
															a__	:							List[Any]      						= ACTaFN[config.hidden_act]
											else:
															a__	:							Union[str, Any]      						= config.hidden_act




							def 		_A	(							self ,					__UpperCAmelCase				):







											"""simple docstring"""



											a__	:							Union[str, Any]      						= self.conva(__UpperCAmelCase				)
											a__	:							Union[str, Any]      						= self.act_fn(__UpperCAmelCase				)
											a__	:							Dict      						= self.drop(__UpperCAmelCase				)
											a__	:							Union[str, Any]      						= self.conva(__UpperCAmelCase				)
											a__	:							int      						= self.drop(__UpperCAmelCase				)

											return hidden_states





class 				_a						(       nn.Module			):





							'''simple docstring'''

							def __init__(							self ,					__UpperCAmelCase ,					__UpperCAmelCase ,					__UpperCAmelCase ,					__UpperCAmelCase ,					__UpperCAmelCase ,					__UpperCAmelCase				):







											"""simple docstring"""



											super().__init__()
											a__	:							Any      						= PoolFormerPooling(__UpperCAmelCase				)
											a__	:							Any      						= PoolFormerOutput(__UpperCAmelCase ,					__UpperCAmelCase ,					__UpperCAmelCase ,					__UpperCAmelCase				)
											a__	:							Any      						= PoolFormerGroupNorm(__UpperCAmelCase				)
											a__	:							Dict      						= PoolFormerGroupNorm(__UpperCAmelCase				)

											# Useful for training neural nets
											a__	:							List[Any]      						= PoolFormerDropPath(__UpperCAmelCase				) if drop_path > 0.0 else nn.Identity()
											a__	:							List[str]      						= config.use_layer_scale
											if config.use_layer_scale:
															a__	:							str      						= nn.Parameter(
															    config.layer_scale_init_value * torch.ones((__UpperCAmelCase)				) ,					requires_grad=__UpperCAmelCase				)
															a__	:							int      						= nn.Parameter(
															    config.layer_scale_init_value * torch.ones((__UpperCAmelCase)				) ,					requires_grad=__UpperCAmelCase				)




							def 		_A	(							self ,					__UpperCAmelCase				):







											"""simple docstring"""



											if self.use_layer_scale:
															a__	:							Any      						= self.pooling(self.before_norm(__UpperCAmelCase				)				)
															a__	:							int      						= self.layer_scale_a.unsqueeze(-1				).unsqueeze(-1				) * pooling_output
															# First residual connection
															a__	:							Optional[int]      						= hidden_states + self.drop_path(__UpperCAmelCase				)
															a__	:							Dict      						= ()

															a__	:							List[Any]      						= self.output(self.after_norm(__UpperCAmelCase				)				)
															a__	:							Union[str, Any]      						= self.layer_scale_a.unsqueeze(-1				).unsqueeze(-1				) * layer_output
															# Second residual connection
															a__	:							Optional[Any]      						= hidden_states + self.drop_path(__UpperCAmelCase				)

															a__	:							Optional[int]      						= (output,) + outputs
															return outputs

											else:
															a__	:							Optional[int]      						= self.drop_path(self.pooling(self.before_norm(__UpperCAmelCase				)				)				)
															# First residual connection
															a__	:							Tuple      						= pooling_output + hidden_states
															a__	:							Tuple      						= ()

															# Second residual connection inside the PoolFormerOutput block
															a__	:							Optional[int]      						= self.drop_path(self.output(self.after_norm(__UpperCAmelCase				)				)				)
															a__	:							str      						= hidden_states + layer_output

															a__	:							Optional[Any]      						= (output,) + outputs
															return outputs





class 				_a						(       nn.Module			):





							'''simple docstring'''

							def __init__(							self ,					__UpperCAmelCase				):







											"""simple docstring"""



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

											# patch embeddings
											a__	:							List[Any]      						= []
											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__	:							Any      						= nn.ModuleList(__UpperCAmelCase				)

											# Transformer blocks
											a__	:							Optional[int]      						= []
											a__	:							List[Any]      						= 0
											for i in range(config.num_encoder_blocks				):
															# each block consists of layers
															a__	:							str      						= []
															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				)				)

											a__	:							Any      						= nn.ModuleList(__UpperCAmelCase				)




							def 		_A	(							self ,					__UpperCAmelCase ,					__UpperCAmelCase=False ,					__UpperCAmelCase=True				):







											"""simple docstring"""



											a__	:							int      						= () if output_hidden_states else None

											a__	:							str      						= pixel_values
											for idx, layers in enumerate(zip(self.patch_embeddings ,					self.block				)				):
															a__						,   a__	:							Optional[Any]      						= layers
															# Get patch embeddings from hidden_states
															a__	:							Any      						= embedding_layer(__UpperCAmelCase				)
															# Send the embeddings through the blocks
															for _, blk in enumerate(__UpperCAmelCase				):
																			a__	:							List[Any]      						= blk(__UpperCAmelCase				)
																			a__	:							Tuple      						= layer_outputs[0]

															if output_hidden_states:
																			a__	:							Optional[Any]      						= 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 				_a						(       SCREAMING_SNAKE_CASE			):





							'''simple docstring'''

							A				:Optional[int] =       PoolFormerConfig
							A				:List[str] =       "poolformer"
							A				:Tuple =       "pixel_values"
							A				:List[str] =       True



							def 		_A	(							self ,					__UpperCAmelCase				):







											"""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 		_A	(							self ,					__UpperCAmelCase ,					__UpperCAmelCase=False				):







											"""simple docstring"""



											if isinstance(__UpperCAmelCase ,					__UpperCAmelCase				):
															a__	:							Dict      						= value





lowerCamelCase		       =      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.
"""

lowerCamelCase		       =      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."       ,  SCREAMING_SNAKE_CASE       ,  )
class 				_a						(       SCREAMING_SNAKE_CASE			):





							'''simple docstring'''

							def __init__(							self ,					__UpperCAmelCase				):







											"""simple docstring"""



											super().__init__(__UpperCAmelCase				)
											a__	:							Optional[Any]      						= config

											a__	:							int      						= PoolFormerEncoder(__UpperCAmelCase				)

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



							def 		_A	(							self				):







											"""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 		_A	(							self ,					__UpperCAmelCase = None ,					__UpperCAmelCase = None ,					__UpperCAmelCase = None ,					):







											"""simple docstring"""



											a__	:							int      						= (
											    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
											)
											a__	:							str      						= 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__	:							List[Any]      						= self.encoder(
											    __UpperCAmelCase ,					output_hidden_states=__UpperCAmelCase ,					return_dict=__UpperCAmelCase ,					)
											a__	:							Any      						= 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 				_a						(       nn.Module			):





							'''simple docstring'''

							def __init__(							self ,					__UpperCAmelCase				):







											"""simple docstring"""



											super().__init__()
											a__	:							Any      						= nn.Linear(config.hidden_size ,					config.hidden_size				)




							def 		_A	(							self ,					__UpperCAmelCase				):







											"""simple docstring"""



											a__	:							Any      						= self.dense(__UpperCAmelCase				)
											return output





@add_start_docstrings(
    "\n    PoolFormer Model transformer with an image classification head on top\n    "       ,  SCREAMING_SNAKE_CASE       ,  )
class 				_a						(       SCREAMING_SNAKE_CASE			):





							'''simple docstring'''

							def __init__(							self ,					__UpperCAmelCase				):







											"""simple docstring"""



											super().__init__(__UpperCAmelCase				)
											a__	:							Optional[Any]      						= config.num_labels
											a__	:							int      						= PoolFormerModel(__UpperCAmelCase				)

											# Final norm
											a__	:							Any      						= PoolFormerGroupNorm(config.hidden_sizes[-1]				)
											# Classifier head
											a__	:							Dict      						= (
											    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 		_A	(							self ,					__UpperCAmelCase = None ,					__UpperCAmelCase = None ,					__UpperCAmelCase = None ,					__UpperCAmelCase = None ,					):







											"""simple docstring"""



											a__	:							str      						= return_dict if return_dict is not None else self.config.use_return_dict

											a__	:							str      						= self.poolformer(
											    __UpperCAmelCase ,					output_hidden_states=__UpperCAmelCase ,					return_dict=__UpperCAmelCase ,					)

											a__	:							Optional[Any]      						= outputs[0]

											a__	:							Union[str, Any]      						= self.classifier(self.norm(__UpperCAmelCase				).mean([-2, -1]				)				)

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

															if self.config.problem_type == "regression":
																			a__	:							int      						= MSELoss()
																			if self.num_labels == 1:
																							a__	:							str      						= loss_fct(logits.squeeze() ,					labels.squeeze()				)
																			else:
																							a__	:							Union[str, Any]      						= loss_fct(__UpperCAmelCase ,					__UpperCAmelCase				)
															elif self.config.problem_type == "single_label_classification":
																			a__	:							Any      						= CrossEntropyLoss()
																			a__	:							Optional[Any]      						= loss_fct(logits.view(-1 ,					self.num_labels				) ,					labels.view(-1				)				)
															elif self.config.problem_type == "multi_label_classification":
																			a__	:							Union[str, Any]      						= BCEWithLogitsLoss()
																			a__	:							Dict      						= loss_fct(__UpperCAmelCase ,					__UpperCAmelCase				)

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

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


