





"""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_     :						Union[str, Any]						=					logging.get_logger(__name__)

# General docstring
lowerCamelCase_     :						Dict						=					'PoolFormerConfig'

# Base docstring
lowerCamelCase_     :						Any						=					'sail/poolformer_s12'
lowerCamelCase_     :						List[str]						=					[1, 5_12, 7, 7]

# Image classification docstring
lowerCamelCase_     :						List[str]						=					'sail/poolformer_s12'
lowerCamelCase_     :						Any						=					'tabby, tabby cat'

lowerCamelCase_     :						Union[str, Any]						=					[
    'sail/poolformer_s12',
    # See all PoolFormer models at https://huggingface.co/models?filter=poolformer
]





def 						UpperCAmelCase__		(     _UpperCAmelCase ,					_UpperCAmelCase = 0.0 ,					_UpperCAmelCase = False					):






  """simple docstring"""





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

class  _UpperCAmelCase  (   nn.Module		):



      '''simple docstring'''

      def __init__(     self					,      snake_case_ = None  ):

        """simple docstring"""

        super().__init__()
        A_  :       Tuple										=  drop_prob

      def 					lowerCamelCase_			(     self					,      snake_case_  ):

        """simple docstring"""

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

      def 					lowerCamelCase_			(     self  ):

        """simple docstring"""

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


class  _UpperCAmelCase  (   nn.Module		):



      '''simple docstring'''

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

        """simple docstring"""

        super().__init__()
        A_  :       Dict										=  patch_size if isinstance(snake_case_					,      collections.abc.Iterable  ) else (patch_size, patch_size)
        A_  :       Tuple										=  stride if isinstance(snake_case_					,      collections.abc.Iterable  ) else (stride, stride)
        A_  :       Dict										=  padding if isinstance(snake_case_					,      collections.abc.Iterable  ) else (padding, padding)

        A_  :       Dict										=  nn.Convad(snake_case_					,      snake_case_					,      kernel_size=snake_case_					,      stride=snake_case_					,      padding=snake_case_  )
        A_  :       Optional[int]										=  norm_layer(snake_case_  ) if norm_layer else nn.Identity()

      def 					lowerCamelCase_			(     self					,      snake_case_  ):

        """simple docstring"""

        A_  :       Union[str, Any]										=  self.projection(snake_case_  )
        A_  :       List[Any]										=  self.norm(snake_case_  )
        return embeddings


class  _UpperCAmelCase  (   nn.GroupNorm		):



      '''simple docstring'''

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

        """simple docstring"""

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


class  _UpperCAmelCase  (   nn.Module		):



      '''simple docstring'''

      def __init__(     self					,      snake_case_  ):

        """simple docstring"""

        super().__init__()
        A_  :       Any										=  nn.AvgPoolad(snake_case_					,      stride=1					,      padding=pool_size // 2					,      count_include_pad=snake_case_  )

      def 					lowerCamelCase_			(     self					,      snake_case_  ):

        """simple docstring"""

        return self.pool(snake_case_  ) - hidden_states


class  _UpperCAmelCase  (   nn.Module		):



      '''simple docstring'''

      def __init__(     self					,      snake_case_					,      snake_case_					,      snake_case_					,      snake_case_  ):

        """simple docstring"""

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

      def 					lowerCamelCase_			(     self					,      snake_case_  ):

        """simple docstring"""

        A_  :       Any										=  self.conva(snake_case_  )
        A_  :       Tuple										=  self.act_fn(snake_case_  )
        A_  :       Optional[int]										=  self.drop(snake_case_  )
        A_  :       List[str]										=  self.conva(snake_case_  )
        A_  :       Union[str, Any]										=  self.drop(snake_case_  )

        return hidden_states


class  _UpperCAmelCase  (   nn.Module		):



      '''simple docstring'''

      def __init__(     self					,      snake_case_					,      snake_case_					,      snake_case_					,      snake_case_					,      snake_case_					,      snake_case_  ):

        """simple docstring"""

        super().__init__()
        A_  :       Any										=  PoolFormerPooling(snake_case_  )
        A_  :       str										=  PoolFormerOutput(snake_case_					,      snake_case_					,      snake_case_					,      snake_case_  )
        A_  :       str										=  PoolFormerGroupNorm(snake_case_  )
        A_  :       List[Any]										=  PoolFormerGroupNorm(snake_case_  )

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

      def 					lowerCamelCase_			(     self					,      snake_case_  ):

        """simple docstring"""

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

          A_  :       List[Any]										=  self.output(self.after_norm(snake_case_  )  )
          A_  :       str										=  self.layer_scale_a.unsqueeze(-1  ).unsqueeze(-1  ) * layer_output
          # Second residual connection
          A_  :       List[Any]										=  hidden_states + self.drop_path(snake_case_  )

          A_  :       Tuple										=  (output,) + outputs
          return outputs

        else:
          A_  :       Tuple										=  self.drop_path(self.pooling(self.before_norm(snake_case_  )  )  )
          # First residual connection
          A_  :       int										=  pooling_output + hidden_states
          A_  :       Any										=  ()

          # Second residual connection inside the PoolFormerOutput block
          A_  :       List[Any]										=  self.drop_path(self.output(self.after_norm(snake_case_  )  )  )
          A_  :       str										=  hidden_states + layer_output

          A_  :       str										=  (output,) + outputs
          return outputs


class  _UpperCAmelCase  (   nn.Module		):



      '''simple docstring'''

      def __init__(     self					,      snake_case_  ):

        """simple docstring"""

        super().__init__()
        A_  :       List[str]										=  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_  :       Optional[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_  :       int										=  nn.ModuleList(snake_case_  )

        # Transformer blocks
        A_  :       Optional[int]										=  []
        A_  :       Tuple										=  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(
                    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_  :       Dict										=  nn.ModuleList(snake_case_  )

      def 					lowerCamelCase_			(     self					,      snake_case_					,      snake_case_=False					,      snake_case_=True  ):

        """simple docstring"""

        A_  :       str										=  () if output_hidden_states else None

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

          if output_hidden_states:
            A_  :       List[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=snake_case_					,      hidden_states=snake_case_  )


class  _UpperCAmelCase  (   UpperCAmelCase__		):



      '''simple docstring'''
      lowercase_      : Any	    =   PoolFormerConfig
      lowercase_      : List[str]	    =   """poolformer"""
      lowercase_      : Union[str, Any]	    =   """pixel_values"""
      lowercase_      : int	    =   True

      def 					lowerCamelCase_			(     self					,      snake_case_  ):

        """simple docstring"""

        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  ):

        """simple docstring"""

        if isinstance(snake_case_					,      snake_case_  ):
          A_  :       Optional[int]										=  value


lowerCamelCase_     :						Optional[Any]						=					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_     :						Any						=					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."""			,   UpperCAmelCase__			,   )
class  _UpperCAmelCase  (   UpperCAmelCase__		):



      '''simple docstring'''

      def __init__(     self					,      snake_case_  ):

        """simple docstring"""

        super().__init__(snake_case_  )
        A_  :       Optional[int]										=  config

        A_  :       List[Any]										=  PoolFormerEncoder(snake_case_  )

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

      def 					lowerCamelCase_			(     self  ):

        """simple docstring"""

        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					,      ):

        """simple docstring"""

        A_  :       Any										=  (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        A_  :       Tuple										=  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(
            snake_case_					,      output_hidden_states=snake_case_					,      return_dict=snake_case_					,      )
        A_  :       int										=  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  _UpperCAmelCase  (   nn.Module		):



      '''simple docstring'''

      def __init__(     self					,      snake_case_  ):

        """simple docstring"""

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

      def 					lowerCamelCase_			(     self					,      snake_case_  ):

        """simple docstring"""

        A_  :       Dict										=  self.dense(snake_case_  )
        return output


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



      '''simple docstring'''

      def __init__(     self					,      snake_case_  ):

        """simple docstring"""

        super().__init__(snake_case_  )
        A_  :       Dict										=  config.num_labels
        A_  :       str										=  PoolFormerModel(snake_case_  )

        # Final norm
        A_  :       Any										=  PoolFormerGroupNorm(config.hidden_sizes[-1]  )
        # Classifier head
        A_  :       List[str]										=  (
            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					,      ):

        """simple docstring"""

        A_  :       Tuple										=  return_dict if return_dict is not None else self.config.use_return_dict

        A_  :       List[Any]										=  self.poolformer(
            snake_case_					,      output_hidden_states=snake_case_					,      return_dict=snake_case_					,      )

        A_  :       Optional[int]										=  outputs[0]

        A_  :       Optional[int]										=  self.classifier(self.norm(snake_case_  ).mean([-2, -1]  )  )

        A_  :       Optional[Any]										=  None
        if labels is not None:
          if self.config.problem_type is None:
            if self.num_labels == 1:
              A_  :       Any										=  'regression'
            elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
              A_  :       Tuple										=  'single_label_classification'
            else:
              A_  :       str										=  'multi_label_classification'

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

        if not return_dict:
          A_  :       Optional[int]										=  (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 torch

from diffusers import DiffusionPipeline

class  _UpperCAmelCase  (   UpperCAmelCase__		):



      '''simple docstring'''

      def __init__(     self					,      snake_case_					,      snake_case_  ):

        """simple docstring"""

        super().__init__()

        self.register_modules(unet=snake_case_					,      scheduler=snake_case_  )

      def __call__(     self  ):

        """simple docstring"""

        A_  :       Optional[Any]										=  torch.randn(
            (1, self.unet.config.in_channels, self.unet.config.sample_size, self.unet.config.sample_size)					,      )
        A_  :       List[str]										=  1

        A_  :       List[str]										=  self.unet(snake_case_					,      snake_case_  ).sample
        A_  :       Optional[int]										=  self.scheduler.step(snake_case_					,      snake_case_					,      snake_case_  ).prev_sample

        A_  :       List[Any]										=  scheduler_output - scheduler_output + torch.ones_like(snake_case_  )

        return result