






from typing import Any, Dict, Optional

import torch
import torch.nn.functional as F
from torch import nn

from ..utils import maybe_allow_in_graph
from .activations import get_activation
from .attention_processor import Attention
from .embeddings import CombinedTimestepLabelEmbeddings





@maybe_allow_in_graph
class 			__a  (    nn.Module       ):







 def __init__(		self					:   int							,   SCREAMING_SNAKE_CASE					:   int							,   SCREAMING_SNAKE_CASE					:   int							,   SCREAMING_SNAKE_CASE					:   int							,   SCREAMING_SNAKE_CASE					:   str=0.0							,   SCREAMING_SNAKE_CASE					:   Optional[int] = None							,   SCREAMING_SNAKE_CASE					:   str = "geglu"							,   SCREAMING_SNAKE_CASE					:   Optional[int] = None							,   SCREAMING_SNAKE_CASE					:   bool = False							,   SCREAMING_SNAKE_CASE					:   bool = False							,   SCREAMING_SNAKE_CASE					:   bool = False							,   SCREAMING_SNAKE_CASE					:   bool = False							,   SCREAMING_SNAKE_CASE					:   bool = True							,   SCREAMING_SNAKE_CASE					:   str = "layer_norm"							,   SCREAMING_SNAKE_CASE					:   bool = False							,   ):


      '''simple docstring'''






      super().__init__()
      UpperCamelCase__			:   Optional[int]       				=  only_cross_attention

      UpperCamelCase__			:   str       				=  (num_embeds_ada_norm is not None) and norm_type == "ada_norm_zero"
      UpperCamelCase__			:   Optional[Any]       				=  (num_embeds_ada_norm is not None) and norm_type == "ada_norm"

      if norm_type in ("ada_norm", "ada_norm_zero") and num_embeds_ada_norm is None:
           raise ValueError(
               F'`norm_type` is set to {norm_type}, but `num_embeds_ada_norm` is not defined. Please make sure to'
               F' define `num_embeds_ada_norm` if setting `norm_type` to {norm_type}.'						)

      # Define 3 blocks. Each block has its own normalization layer.
      # 1. Self-Attn
      if self.use_ada_layer_norm:
           UpperCamelCase__			:   List[Any]       				=  AdaLayerNorm(SCREAMING_SNAKE_CASE							,   SCREAMING_SNAKE_CASE						)
      elif self.use_ada_layer_norm_zero:
           UpperCamelCase__			:   Dict       				=  AdaLayerNormZero(SCREAMING_SNAKE_CASE							,   SCREAMING_SNAKE_CASE						)
      else:
           UpperCamelCase__			:   str       				=  nn.LayerNorm(SCREAMING_SNAKE_CASE							,   elementwise_affine=SCREAMING_SNAKE_CASE						)
      UpperCamelCase__			:   List[Any]       				=  Attention(
          query_dim=SCREAMING_SNAKE_CASE							,   heads=SCREAMING_SNAKE_CASE							,   dim_head=SCREAMING_SNAKE_CASE							,   dropout=SCREAMING_SNAKE_CASE							,   bias=SCREAMING_SNAKE_CASE							,   cross_attention_dim=cross_attention_dim if only_cross_attention else None							,   upcast_attention=SCREAMING_SNAKE_CASE							,   )

      # 2. Cross-Attn
      if cross_attention_dim is not None or double_self_attention:
           # We currently only use AdaLayerNormZero for self attention where there will only be one attention block.
           # I.e. the number of returned modulation chunks from AdaLayerZero would not make sense if returned during
           # the second cross attention block.
           UpperCamelCase__			:   Optional[int]       				=  (
               AdaLayerNorm(SCREAMING_SNAKE_CASE							,   SCREAMING_SNAKE_CASE						)
               if self.use_ada_layer_norm
               else nn.LayerNorm(SCREAMING_SNAKE_CASE							,   elementwise_affine=SCREAMING_SNAKE_CASE						)
           )
           UpperCamelCase__			:   int       				=  Attention(
               query_dim=SCREAMING_SNAKE_CASE							,   cross_attention_dim=cross_attention_dim if not double_self_attention else None							,   heads=SCREAMING_SNAKE_CASE							,   dim_head=SCREAMING_SNAKE_CASE							,   dropout=SCREAMING_SNAKE_CASE							,   bias=SCREAMING_SNAKE_CASE							,   upcast_attention=SCREAMING_SNAKE_CASE							,   )  # is self-attn if encoder_hidden_states is none
      else:
           UpperCamelCase__			:   Union[str, Any]       				=  None
           UpperCamelCase__			:   str       				=  None

      # 3. Feed-forward
      UpperCamelCase__			:   List[Any]       				=  nn.LayerNorm(SCREAMING_SNAKE_CASE							,   elementwise_affine=SCREAMING_SNAKE_CASE						)
      UpperCamelCase__			:   Optional[Any]       				=  FeedForward(SCREAMING_SNAKE_CASE							,   dropout=SCREAMING_SNAKE_CASE							,   activation_fn=SCREAMING_SNAKE_CASE							,   final_dropout=SCREAMING_SNAKE_CASE						)

      # let chunk size default to None
      UpperCamelCase__			:   List[str]       				=  None
      UpperCamelCase__			:   Optional[Any]       				=  0







 def 					__lowercase					(		self					:   List[Any]							,   SCREAMING_SNAKE_CASE					:   Optional[int]							,   SCREAMING_SNAKE_CASE					:   int						):


      '''simple docstring'''






      UpperCamelCase__			:   int       				=  chunk_size
      UpperCamelCase__			:   List[str]       				=  dim







 def 					__lowercase					(		self					:   Dict							,   SCREAMING_SNAKE_CASE					:   torch.FloatTensor							,   SCREAMING_SNAKE_CASE					:   Optional[torch.FloatTensor] = None							,   SCREAMING_SNAKE_CASE					:   Optional[torch.FloatTensor] = None							,   SCREAMING_SNAKE_CASE					:   Optional[torch.FloatTensor] = None							,   SCREAMING_SNAKE_CASE					:   Optional[torch.LongTensor] = None							,   SCREAMING_SNAKE_CASE					:   Dict[str, Any] = None							,   SCREAMING_SNAKE_CASE					:   Optional[torch.LongTensor] = None							,   ):


      '''simple docstring'''






      if self.use_ada_layer_norm:
           UpperCamelCase__			:   str       				=  self.norma(SCREAMING_SNAKE_CASE							,   SCREAMING_SNAKE_CASE						)
      elif self.use_ada_layer_norm_zero:
           UpperCamelCase__				,     UpperCamelCase__				,     UpperCamelCase__				,     UpperCamelCase__				,     UpperCamelCase__			:   Any       				=  self.norma(
               SCREAMING_SNAKE_CASE							,   SCREAMING_SNAKE_CASE							,   SCREAMING_SNAKE_CASE							,   hidden_dtype=hidden_states.dtype						)
      else:
           UpperCamelCase__			:   List[Any]       				=  self.norma(SCREAMING_SNAKE_CASE						)

      UpperCamelCase__			:   int       				=  cross_attention_kwargs if cross_attention_kwargs is not None else {}

      UpperCamelCase__			:   int       				=  self.attna(
          SCREAMING_SNAKE_CASE							,   encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None							,   attention_mask=SCREAMING_SNAKE_CASE							,   **SCREAMING_SNAKE_CASE							,   )
      if self.use_ada_layer_norm_zero:
           UpperCamelCase__			:   int       				=  gate_msa.unsqueeze(1						) * attn_output
      UpperCamelCase__			:   Dict       				=  attn_output + hidden_states

      # 2. Cross-Attention
      if self.attna is not None:
           UpperCamelCase__			:   Optional[int]       				=  (
               self.norma(SCREAMING_SNAKE_CASE							,   SCREAMING_SNAKE_CASE						) if self.use_ada_layer_norm else self.norma(SCREAMING_SNAKE_CASE						)
           )

           UpperCamelCase__			:   Optional[Any]       				=  self.attna(
               SCREAMING_SNAKE_CASE							,   encoder_hidden_states=SCREAMING_SNAKE_CASE							,   attention_mask=SCREAMING_SNAKE_CASE							,   **SCREAMING_SNAKE_CASE							,   )
           UpperCamelCase__			:   Union[str, Any]       				=  attn_output + hidden_states

      # 3. Feed-forward
      UpperCamelCase__			:   Dict       				=  self.norma(SCREAMING_SNAKE_CASE						)

      if self.use_ada_layer_norm_zero:
           UpperCamelCase__			:   Optional[Any]       				=  norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]

      if self._chunk_size is not None:
           # "feed_forward_chunk_size" can be used to save memory
           if norm_hidden_states.shape[self._chunk_dim] % self._chunk_size != 0:
                raise ValueError(
                    F'`hidden_states` dimension to be chunked: {norm_hidden_states.shape[self._chunk_dim]} has to be divisible by chunk size: {self._chunk_size}. Make sure to set an appropriate `chunk_size` when calling `unet.enable_forward_chunking`.'						)

           UpperCamelCase__			:   Union[str, Any]       				=  norm_hidden_states.shape[self._chunk_dim] // self._chunk_size
           UpperCamelCase__			:   List[str]       				=  torch.cat(
               [self.ff(SCREAMING_SNAKE_CASE						) for hid_slice in norm_hidden_states.chunk(SCREAMING_SNAKE_CASE							,   dim=self._chunk_dim						)]							,   dim=self._chunk_dim							,   )
      else:
           UpperCamelCase__			:   Union[str, Any]       				=  self.ff(SCREAMING_SNAKE_CASE						)

      if self.use_ada_layer_norm_zero:
           UpperCamelCase__			:   List[str]       				=  gate_mlp.unsqueeze(1						) * ff_output

      UpperCamelCase__			:   str       				=  ff_output + hidden_states

      return hidden_states







class 			__a  (    nn.Module       ):







 def __init__(		self					:   Dict							,   SCREAMING_SNAKE_CASE					:   int							,   SCREAMING_SNAKE_CASE					:   Optional[int] = None							,   SCREAMING_SNAKE_CASE					:   int = 4							,   SCREAMING_SNAKE_CASE					:   float = 0.0							,   SCREAMING_SNAKE_CASE					:   str = "geglu"							,   SCREAMING_SNAKE_CASE					:   bool = False							,   ):


      '''simple docstring'''






      super().__init__()
      UpperCamelCase__			:   Tuple       				=  int(dim * mult						)
      UpperCamelCase__			:   Dict       				=  dim_out if dim_out is not None else dim

      if activation_fn == "gelu":
           UpperCamelCase__			:   List[str]       				=  GELU(SCREAMING_SNAKE_CASE							,   SCREAMING_SNAKE_CASE						)
      if activation_fn == "gelu-approximate":
           UpperCamelCase__			:   Dict       				=  GELU(SCREAMING_SNAKE_CASE							,   SCREAMING_SNAKE_CASE							,   approximate="tanh"						)
      elif activation_fn == "geglu":
           UpperCamelCase__			:   str       				=  GEGLU(SCREAMING_SNAKE_CASE							,   SCREAMING_SNAKE_CASE						)
      elif activation_fn == "geglu-approximate":
           UpperCamelCase__			:   str       				=  ApproximateGELU(SCREAMING_SNAKE_CASE							,   SCREAMING_SNAKE_CASE						)

      UpperCamelCase__			:   Any       				=  nn.ModuleList([]						)
      # project in
      self.net.append(SCREAMING_SNAKE_CASE						)
      # project dropout
      self.net.append(nn.Dropout(SCREAMING_SNAKE_CASE						)						)
      # project out
      self.net.append(nn.Linear(SCREAMING_SNAKE_CASE							,   SCREAMING_SNAKE_CASE						)						)
      # FF as used in Vision Transformer, MLP-Mixer, etc. have a final dropout
      if final_dropout:
           self.net.append(nn.Dropout(SCREAMING_SNAKE_CASE						)						)







 def 					__lowercase					(		self					:   Dict							,   SCREAMING_SNAKE_CASE					:   Any						):


      '''simple docstring'''






      for module in self.net:
           UpperCamelCase__			:   Optional[Any]       				=  module(SCREAMING_SNAKE_CASE						)
      return hidden_states







class 			__a  (    nn.Module       ):







 def __init__(		self					:   int							,   SCREAMING_SNAKE_CASE					:   int							,   SCREAMING_SNAKE_CASE					:   int							,   SCREAMING_SNAKE_CASE					:   str = "none"						):


      '''simple docstring'''






      super().__init__()
      UpperCamelCase__			:   Optional[int]       				=  nn.Linear(SCREAMING_SNAKE_CASE							,   SCREAMING_SNAKE_CASE						)
      UpperCamelCase__			:   Any       				=  approximate







 def 					__lowercase					(		self					:   Optional[Any]							,   SCREAMING_SNAKE_CASE					:   str						):


      '''simple docstring'''






      if gate.device.type != "mps":
           return F.gelu(SCREAMING_SNAKE_CASE							,   approximate=self.approximate						)
      # mps: gelu is not implemented for float16
      return F.gelu(gate.to(dtype=torch.floataa						)							,   approximate=self.approximate						).to(dtype=gate.dtype						)







 def 					__lowercase					(		self					:   str							,   SCREAMING_SNAKE_CASE					:   List[Any]						):


      '''simple docstring'''






      UpperCamelCase__			:   Tuple       				=  self.proj(SCREAMING_SNAKE_CASE						)
      UpperCamelCase__			:   Optional[int]       				=  self.gelu(SCREAMING_SNAKE_CASE						)
      return hidden_states







class 			__a  (    nn.Module       ):







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


      '''simple docstring'''






      super().__init__()
      UpperCamelCase__			:   int       				=  nn.Linear(SCREAMING_SNAKE_CASE							,   dim_out * 2						)







 def 					__lowercase					(		self					:   List[str]							,   SCREAMING_SNAKE_CASE					:   Tuple						):


      '''simple docstring'''






      if gate.device.type != "mps":
           return F.gelu(SCREAMING_SNAKE_CASE						)
      # mps: gelu is not implemented for float16
      return F.gelu(gate.to(dtype=torch.floataa						)						).to(dtype=gate.dtype						)







 def 					__lowercase					(		self					:   Optional[int]							,   SCREAMING_SNAKE_CASE					:   Optional[int]						):


      '''simple docstring'''






      UpperCamelCase__				,     UpperCamelCase__			:   List[str]       				=  self.proj(SCREAMING_SNAKE_CASE						).chunk(2							,   dim=-1						)
      return hidden_states * self.gelu(SCREAMING_SNAKE_CASE						)







class 			__a  (    nn.Module       ):







 def __init__(		self					:   Any							,   SCREAMING_SNAKE_CASE					:   int							,   SCREAMING_SNAKE_CASE					:   int						):


      '''simple docstring'''






      super().__init__()
      UpperCamelCase__			:   int       				=  nn.Linear(SCREAMING_SNAKE_CASE							,   SCREAMING_SNAKE_CASE						)







 def 					__lowercase					(		self					:   Optional[Any]							,   SCREAMING_SNAKE_CASE					:   Union[str, Any]						):


      '''simple docstring'''






      UpperCamelCase__			:   Tuple       				=  self.proj(SCREAMING_SNAKE_CASE						)
      return x * torch.sigmoid(1.7_0_2 * x						)







class 			__a  (    nn.Module       ):







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


      '''simple docstring'''






      super().__init__()
      UpperCamelCase__			:   Optional[int]       				=  nn.Embedding(SCREAMING_SNAKE_CASE							,   SCREAMING_SNAKE_CASE						)
      UpperCamelCase__			:   Any       				=  nn.SiLU()
      UpperCamelCase__			:   Optional[Any]       				=  nn.Linear(SCREAMING_SNAKE_CASE							,   embedding_dim * 2						)
      UpperCamelCase__			:   Union[str, Any]       				=  nn.LayerNorm(SCREAMING_SNAKE_CASE							,   elementwise_affine=SCREAMING_SNAKE_CASE						)







 def 					__lowercase					(		self					:   Union[str, Any]							,   SCREAMING_SNAKE_CASE					:   Optional[int]							,   SCREAMING_SNAKE_CASE					:   int						):


      '''simple docstring'''






      UpperCamelCase__			:   Dict       				=  self.linear(self.silu(self.emb(SCREAMING_SNAKE_CASE						)						)						)
      UpperCamelCase__				,     UpperCamelCase__			:   Optional[Any]       				=  torch.chunk(SCREAMING_SNAKE_CASE							,   2						)
      UpperCamelCase__			:   Optional[Any]       				=  self.norm(SCREAMING_SNAKE_CASE						) * (1 + scale) + shift
      return x







class 			__a  (    nn.Module       ):







 def __init__(		self					:   str							,   SCREAMING_SNAKE_CASE					:   Optional[int]							,   SCREAMING_SNAKE_CASE					:   List[Any]						):


      '''simple docstring'''






      super().__init__()

      UpperCamelCase__			:   Tuple       				=  CombinedTimestepLabelEmbeddings(SCREAMING_SNAKE_CASE							,   SCREAMING_SNAKE_CASE						)

      UpperCamelCase__			:   List[Any]       				=  nn.SiLU()
      UpperCamelCase__			:   str       				=  nn.Linear(SCREAMING_SNAKE_CASE							,   6 * embedding_dim							,   bias=SCREAMING_SNAKE_CASE						)
      UpperCamelCase__			:   List[Any]       				=  nn.LayerNorm(SCREAMING_SNAKE_CASE							,   elementwise_affine=SCREAMING_SNAKE_CASE							,   eps=1e-6						)







 def 					__lowercase					(		self					:   Dict							,   SCREAMING_SNAKE_CASE					:   int							,   SCREAMING_SNAKE_CASE					:   Dict							,   SCREAMING_SNAKE_CASE					:   List[str]							,   SCREAMING_SNAKE_CASE					:   List[Any]=None						):


      '''simple docstring'''






      UpperCamelCase__			:   List[str]       				=  self.linear(self.silu(self.emb(SCREAMING_SNAKE_CASE							,   SCREAMING_SNAKE_CASE							,   hidden_dtype=SCREAMING_SNAKE_CASE						)						)						)
      UpperCamelCase__				,     UpperCamelCase__				,     UpperCamelCase__				,     UpperCamelCase__				,     UpperCamelCase__				,     UpperCamelCase__			:   str       				=  emb.chunk(6							,   dim=1						)
      UpperCamelCase__			:   str       				=  self.norm(SCREAMING_SNAKE_CASE						) * (1 + scale_msa[:, None]) + shift_msa[:, None]
      return x, gate_msa, shift_mlp, scale_mlp, gate_mlp







class 			__a  (    nn.Module       ):







 def __init__(		self					:   Any							,   SCREAMING_SNAKE_CASE					:   int							,   SCREAMING_SNAKE_CASE					:   int							,   SCREAMING_SNAKE_CASE					:   int							,   SCREAMING_SNAKE_CASE					:   Optional[str] = None							,   SCREAMING_SNAKE_CASE					:   float = 1e-5						):


      '''simple docstring'''






      super().__init__()
      UpperCamelCase__			:   Optional[Any]       				=  num_groups
      UpperCamelCase__			:   Tuple       				=  eps

      if act_fn is None:
           UpperCamelCase__			:   Optional[Any]       				=  None
      else:
           UpperCamelCase__			:   Optional[int]       				=  get_activation(SCREAMING_SNAKE_CASE						)

      UpperCamelCase__			:   Optional[Any]       				=  nn.Linear(SCREAMING_SNAKE_CASE							,   out_dim * 2						)







 def 					__lowercase					(		self					:   Any							,   SCREAMING_SNAKE_CASE					:   str							,   SCREAMING_SNAKE_CASE					:   List[str]						):


      '''simple docstring'''






      if self.act:
           UpperCamelCase__			:   Optional[int]       				=  self.act(SCREAMING_SNAKE_CASE						)
      UpperCamelCase__			:   Union[str, Any]       				=  self.linear(SCREAMING_SNAKE_CASE						)
      UpperCamelCase__			:   Any       				=  emb[:, :, None, None]
      UpperCamelCase__				,     UpperCamelCase__			:   Optional[int]       				=  emb.chunk(2							,   dim=1						)

      UpperCamelCase__			:   Optional[Any]       				=  F.group_norm(SCREAMING_SNAKE_CASE							,   self.num_groups							,   eps=self.eps						)
      UpperCamelCase__			:   Union[str, Any]       				=  x * (1 + scale) + shift
      return x







from typing import Dict, List, Optional, Tuple, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import (
    center_crop,
    flip_channel_order,
    get_resize_output_image_size,
    rescale,
    resize,
    to_channel_dimension_format,
)
from ...image_utils import (
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_torch_available, is_torch_tensor, is_vision_available, logging


if is_vision_available():
 import PIL

if is_torch_available():
 import torch


lowerCamelCase	:							Any				=logging.get_logger(__name__)





class 			__a  (    A__       ):
 _lowerCAmelCase       :       List[str]        = ['''pixel_values''']







 def __init__(		self					:   List[str]							,   SCREAMING_SNAKE_CASE					:   bool = True							,   SCREAMING_SNAKE_CASE					:   Dict[str, int] = None							,   SCREAMING_SNAKE_CASE					:   PILImageResampling = PILImageResampling.BILINEAR							,   SCREAMING_SNAKE_CASE					:   bool = True							,   SCREAMING_SNAKE_CASE					:   Union[int, float] = 1 / 2_55							,   SCREAMING_SNAKE_CASE					:   bool = True							,   SCREAMING_SNAKE_CASE					:   Dict[str, int] = None							,   SCREAMING_SNAKE_CASE					:   bool = True							,   **SCREAMING_SNAKE_CASE					:   Tuple							,   ):


      '''simple docstring'''






      super().__init__(**SCREAMING_SNAKE_CASE						)
      UpperCamelCase__			:   List[str]       				=  size if size is not None else {"shortest_edge": 2_24}
      UpperCamelCase__			:   Any       				=  get_size_dict(SCREAMING_SNAKE_CASE							,   default_to_square=SCREAMING_SNAKE_CASE						)
      UpperCamelCase__			:   List[Any]       				=  crop_size if crop_size is not None else {"height": 2_56, "width": 2_56}
      UpperCamelCase__			:   int       				=  get_size_dict(SCREAMING_SNAKE_CASE							,   param_name="crop_size"						)

      UpperCamelCase__			:   Dict       				=  do_resize
      UpperCamelCase__			:   List[str]       				=  size
      UpperCamelCase__			:   int       				=  resample
      UpperCamelCase__			:   Optional[int]       				=  do_rescale
      UpperCamelCase__			:   List[Any]       				=  rescale_factor
      UpperCamelCase__			:   Union[str, Any]       				=  do_center_crop
      UpperCamelCase__			:   int       				=  crop_size
      UpperCamelCase__			:   Optional[int]       				=  do_flip_channel_order







 def 					__lowercase					(		self					:   Optional[Any]							,   SCREAMING_SNAKE_CASE					:   np.ndarray							,   SCREAMING_SNAKE_CASE					:   Dict[str, int]							,   SCREAMING_SNAKE_CASE					:   PILImageResampling = PIL.Image.BILINEAR							,   SCREAMING_SNAKE_CASE					:   Optional[Union[str, ChannelDimension]] = None							,   **SCREAMING_SNAKE_CASE					:   Optional[Any]							,   ):


      '''simple docstring'''






      UpperCamelCase__			:   Union[str, Any]       				=  get_size_dict(SCREAMING_SNAKE_CASE							,   default_to_square=SCREAMING_SNAKE_CASE						)
      if "shortest_edge" not in size:
           raise ValueError(F'The `size` dictionary must contain the key `shortest_edge`. Got {size.keys()}'						)
      UpperCamelCase__			:   Any       				=  get_resize_output_image_size(SCREAMING_SNAKE_CASE							,   size=size["shortest_edge"]							,   default_to_square=SCREAMING_SNAKE_CASE						)
      return resize(SCREAMING_SNAKE_CASE							,   size=SCREAMING_SNAKE_CASE							,   resample=SCREAMING_SNAKE_CASE							,   data_format=SCREAMING_SNAKE_CASE							,   **SCREAMING_SNAKE_CASE						)







 def 					__lowercase					(		self					:   Optional[int]							,   SCREAMING_SNAKE_CASE					:   np.ndarray							,   SCREAMING_SNAKE_CASE					:   Dict[str, int]							,   SCREAMING_SNAKE_CASE					:   Optional[Union[str, ChannelDimension]] = None							,   **SCREAMING_SNAKE_CASE					:   Optional[Any]							,   ):


      '''simple docstring'''






      UpperCamelCase__			:   List[Any]       				=  get_size_dict(SCREAMING_SNAKE_CASE						)
      if "height" not in size or "width" not in size:
           raise ValueError(F'The `size` dictionary must contain the keys `height` and `width`. Got {size.keys()}'						)
      return center_crop(SCREAMING_SNAKE_CASE							,   size=(size["height"], size["width"])							,   data_format=SCREAMING_SNAKE_CASE							,   **SCREAMING_SNAKE_CASE						)







 def 					__lowercase					(		self					:   int							,   SCREAMING_SNAKE_CASE					:   np.ndarray							,   SCREAMING_SNAKE_CASE					:   Union[int, float]							,   SCREAMING_SNAKE_CASE					:   Optional[Union[str, ChannelDimension]] = None							,   **SCREAMING_SNAKE_CASE					:   Union[str, Any]							,   ):


      '''simple docstring'''






      return rescale(SCREAMING_SNAKE_CASE							,   scale=SCREAMING_SNAKE_CASE							,   data_format=SCREAMING_SNAKE_CASE							,   **SCREAMING_SNAKE_CASE						)







 def 					__lowercase					(		self					:   Optional[int]							,   SCREAMING_SNAKE_CASE					:   np.ndarray							,   SCREAMING_SNAKE_CASE					:   Optional[Union[str, ChannelDimension]] = None						):


      '''simple docstring'''






      return flip_channel_order(SCREAMING_SNAKE_CASE							,   data_format=SCREAMING_SNAKE_CASE						)







 def 					__lowercase					(		self					:   Optional[Any]							,   SCREAMING_SNAKE_CASE					:   ImageInput							,   SCREAMING_SNAKE_CASE					:   bool = None							,   SCREAMING_SNAKE_CASE					:   Dict[str, int] = None							,   SCREAMING_SNAKE_CASE					:   PILImageResampling = None							,   SCREAMING_SNAKE_CASE					:   bool = None							,   SCREAMING_SNAKE_CASE					:   float = None							,   SCREAMING_SNAKE_CASE					:   bool = None							,   SCREAMING_SNAKE_CASE					:   Dict[str, int] = None							,   SCREAMING_SNAKE_CASE					:   bool = None							,   SCREAMING_SNAKE_CASE					:   Optional[Union[str, TensorType]] = None							,   SCREAMING_SNAKE_CASE					:   ChannelDimension = ChannelDimension.FIRST							,   **SCREAMING_SNAKE_CASE					:   Dict							,   ):


      '''simple docstring'''






      UpperCamelCase__			:   List[str]       				=  do_resize if do_resize is not None else self.do_resize
      UpperCamelCase__			:   List[str]       				=  resample if resample is not None else self.resample
      UpperCamelCase__			:   Any       				=  do_rescale if do_rescale is not None else self.do_rescale
      UpperCamelCase__			:   Tuple       				=  rescale_factor if rescale_factor is not None else self.rescale_factor
      UpperCamelCase__			:   Dict       				=  do_center_crop if do_center_crop is not None else self.do_center_crop
      UpperCamelCase__			:   Any       				=  (
          do_flip_channel_order if do_flip_channel_order is not None else self.do_flip_channel_order
      )

      UpperCamelCase__			:   Optional[int]       				=  size if size is not None else self.size
      UpperCamelCase__			:   List[str]       				=  get_size_dict(SCREAMING_SNAKE_CASE							,   default_to_square=SCREAMING_SNAKE_CASE						)
      UpperCamelCase__			:   Any       				=  crop_size if crop_size is not None else self.crop_size
      UpperCamelCase__			:   Union[str, Any]       				=  get_size_dict(SCREAMING_SNAKE_CASE							,   param_name="crop_size"						)

      UpperCamelCase__			:   int       				=  make_list_of_images(SCREAMING_SNAKE_CASE						)

      if not valid_images(SCREAMING_SNAKE_CASE						):
           raise ValueError(
               "Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, "
               "torch.Tensor, tf.Tensor or jax.ndarray."						)

      if do_resize and size is None:
           raise ValueError("Size must be specified if do_resize is True."						)

      if do_rescale and rescale_factor is None:
           raise ValueError("Rescale factor must be specified if do_rescale is True."						)

      if do_center_crop and crop_size is None:
           raise ValueError("Crop size must be specified if do_center_crop is True."						)

      # All transformations expect numpy arrays.
      UpperCamelCase__			:   Tuple       				=  [to_numpy_array(SCREAMING_SNAKE_CASE						) for image in images]

      if do_resize:
           UpperCamelCase__			:   Optional[int]       				=  [self.resize(image=SCREAMING_SNAKE_CASE							,   size=SCREAMING_SNAKE_CASE							,   resample=SCREAMING_SNAKE_CASE						) for image in images]

      if do_center_crop:
           UpperCamelCase__			:   Any       				=  [self.center_crop(image=SCREAMING_SNAKE_CASE							,   size=SCREAMING_SNAKE_CASE						) for image in images]

      if do_rescale:
           UpperCamelCase__			:   str       				=  [self.rescale(image=SCREAMING_SNAKE_CASE							,   scale=SCREAMING_SNAKE_CASE						) for image in images]

      # the pretrained checkpoints assume images are BGR, not RGB
      if do_flip_channel_order:
           UpperCamelCase__			:   Any       				=  [self.flip_channel_order(image=SCREAMING_SNAKE_CASE						) for image in images]

      UpperCamelCase__			:   Optional[int]       				=  [to_channel_dimension_format(SCREAMING_SNAKE_CASE							,   SCREAMING_SNAKE_CASE						) for image in images]

      UpperCamelCase__			:   Optional[int]       				=  {"pixel_values": images}
      return BatchFeature(data=SCREAMING_SNAKE_CASE							,   tensor_type=SCREAMING_SNAKE_CASE						)







 def 					__lowercase					(		self					:   Dict							,   SCREAMING_SNAKE_CASE					:   List[Any]							,   SCREAMING_SNAKE_CASE					:   List[Tuple] = None						):


      '''simple docstring'''






      UpperCamelCase__			:   Optional[int]       				=  outputs.logits

      # Resize logits and compute semantic segmentation maps
      if target_sizes is not None:
           if len(SCREAMING_SNAKE_CASE						) != len(SCREAMING_SNAKE_CASE						):
                raise ValueError(
                    "Make sure that you pass in as many target sizes as the batch dimension of the logits"						)

           if is_torch_tensor(SCREAMING_SNAKE_CASE						):
                UpperCamelCase__			:   Optional[Any]       				=  target_sizes.numpy()

           UpperCamelCase__			:   Any       				=  []

           for idx in range(len(SCREAMING_SNAKE_CASE						)						):
                UpperCamelCase__			:   Optional[Any]       				=  torch.nn.functional.interpolate(
                    logits[idx].unsqueeze(dim=0						)							,   size=target_sizes[idx]							,   mode="bilinear"							,   align_corners=SCREAMING_SNAKE_CASE						)
                UpperCamelCase__			:   List[str]       				=  resized_logits[0].argmax(dim=0						)
                semantic_segmentation.append(SCREAMING_SNAKE_CASE						)
      else:
           UpperCamelCase__			:   List[str]       				=  logits.argmax(dim=1						)
           UpperCamelCase__			:   Dict       				=  [semantic_segmentation[i] for i in range(semantic_segmentation.shape[0]						)]

      return semantic_segmentation