

from dataclasses import dataclass
from typing import Optional

import numpy as np
import torch
import torch.nn as nn

from ..utils import BaseOutput, is_torch_version, randn_tensor
from .attention_processor import SpatialNorm
from .unet_ad_blocks import UNetMidBlockaD, get_down_block, get_up_block





@dataclass
class A__    ( UpperCamelCase__ ):



 """simple docstring"""







 _lowercase     :					List[Any]      =						42





class A__    ( nn.Module ):



 """simple docstring"""







 def __init__(		self :	Optional[int]     ,						A_ :	int=3     ,						A_ :	Union[str, Any]=3     ,						A_ :	Any=("DownEncoderBlock2D",)     ,						A_ :	Tuple=(6_4,)     ,						A_ :	str=2     ,						A_ :	List[Any]=3_2     ,						A_ :	Tuple="silu"     ,						A_ :	Dict=True     ,						):






      '''simple docstring'''




      super().__init__()
      _lowerCAmelCase   :					Optional[int]        =    layers_per_block

      _lowerCAmelCase   :					Tuple        =    torch.nn.Convad(
          _a     ,						block_out_channels[0]     ,						kernel_size=3     ,						stride=1     ,						padding=1     ,						)

      _lowerCAmelCase   :					str        =    None
      _lowerCAmelCase   :					int        =    nn.ModuleList([]       )

      # down
      _lowerCAmelCase   :					Tuple        =    block_out_channels[0]
      for i, down_block_type in enumerate(_a       ):
           _lowerCAmelCase   :					Optional[int]        =    output_channel
           _lowerCAmelCase   :					List[Any]        =    block_out_channels[i]
           _lowerCAmelCase   :					List[Any]        =    i == len(_a       ) - 1

           _lowerCAmelCase   :					Union[str, Any]        =    get_down_block(
               _a     ,						num_layers=self.layers_per_block     ,						in_channels=_a     ,						out_channels=_a     ,						add_downsample=not is_final_block     ,						resnet_eps=1E-6     ,						downsample_padding=0     ,						resnet_act_fn=_a     ,						resnet_groups=_a     ,						attention_head_dim=_a     ,						temb_channels=_a     ,						)
           self.down_blocks.append(_a       )

      # mid
      _lowerCAmelCase   :					List[str]        =    UNetMidBlockaD(
          in_channels=block_out_channels[-1]     ,						resnet_eps=1E-6     ,						resnet_act_fn=_a     ,						output_scale_factor=1     ,						resnet_time_scale_shift="default"     ,						attention_head_dim=block_out_channels[-1]     ,						resnet_groups=_a     ,						temb_channels=_a     ,						)

      # out
      _lowerCAmelCase   :					Dict        =    nn.GroupNorm(num_channels=block_out_channels[-1]     ,						num_groups=_a     ,						eps=1E-6       )
      _lowerCAmelCase   :					Optional[int]        =    nn.SiLU()

      _lowerCAmelCase   :					List[str]        =    2 * out_channels if double_z else out_channels
      _lowerCAmelCase   :					List[str]        =    nn.Convad(block_out_channels[-1]     ,						_a     ,						3     ,						padding=1       )

      _lowerCAmelCase   :					int        =    False


 def    __magic_name__       (		self :	Any     ,						A_ :	int       ):






      '''simple docstring'''




      _lowerCAmelCase   :					Dict        =    x
      _lowerCAmelCase   :					Tuple        =    self.conv_in(_a       )

      if self.training and self.gradient_checkpointing:

           def create_custom_forward(A_ :	Tuple       ):
                def custom_forward(*A_ :	Dict       ):
                     return module(*_a       )

                return custom_forward

           # down
           if is_torch_version(">="     ,						"1.11.0"       ):
                for down_block in self.down_blocks:
                     _lowerCAmelCase   :					Tuple        =    torch.utils.checkpoint.checkpoint(
                         create_custom_forward(_a       )     ,						_a     ,						use_reentrant=_a       )
                # middle
                _lowerCAmelCase   :					Optional[int]        =    torch.utils.checkpoint.checkpoint(
                    create_custom_forward(self.mid_block       )     ,						_a     ,						use_reentrant=_a       )
           else:
                for down_block in self.down_blocks:
                     _lowerCAmelCase   :					Optional[Any]        =    torch.utils.checkpoint.checkpoint(create_custom_forward(_a       )     ,						_a       )
                # middle
                _lowerCAmelCase   :					Union[str, Any]        =    torch.utils.checkpoint.checkpoint(create_custom_forward(self.mid_block       )     ,						_a       )

      else:
           # down
           for down_block in self.down_blocks:
                _lowerCAmelCase   :					Optional[int]        =    down_block(_a       )

           # middle
           _lowerCAmelCase   :					int        =    self.mid_block(_a       )

      # post-process
      _lowerCAmelCase   :					Tuple        =    self.conv_norm_out(_a       )
      _lowerCAmelCase   :					Optional[int]        =    self.conv_act(_a       )
      _lowerCAmelCase   :					Optional[int]        =    self.conv_out(_a       )

      return sample





class A__    ( nn.Module ):



 """simple docstring"""







 def __init__(		self :	List[str]     ,						A_ :	int=3     ,						A_ :	List[str]=3     ,						A_ :	Tuple=("UpDecoderBlock2D",)     ,						A_ :	Tuple=(6_4,)     ,						A_ :	Tuple=2     ,						A_ :	str=3_2     ,						A_ :	Optional[int]="silu"     ,						A_ :	Any="group"     ,						):






      '''simple docstring'''




      super().__init__()
      _lowerCAmelCase   :					str        =    layers_per_block

      _lowerCAmelCase   :					str        =    nn.Convad(
          _a     ,						block_out_channels[-1]     ,						kernel_size=3     ,						stride=1     ,						padding=1     ,						)

      _lowerCAmelCase   :					Optional[int]        =    None
      _lowerCAmelCase   :					int        =    nn.ModuleList([]       )

      _lowerCAmelCase   :					Dict        =    in_channels if norm_type == """spatial""" else None

      # mid
      _lowerCAmelCase   :					Tuple        =    UNetMidBlockaD(
          in_channels=block_out_channels[-1]     ,						resnet_eps=1E-6     ,						resnet_act_fn=_a     ,						output_scale_factor=1     ,						resnet_time_scale_shift="default" if norm_type == "group" else norm_type     ,						attention_head_dim=block_out_channels[-1]     ,						resnet_groups=_a     ,						temb_channels=_a     ,						)

      # up
      _lowerCAmelCase   :					Union[str, Any]        =    list(reversed(_a       )       )
      _lowerCAmelCase   :					int        =    reversed_block_out_channels[0]
      for i, up_block_type in enumerate(_a       ):
           _lowerCAmelCase   :					List[Any]        =    output_channel
           _lowerCAmelCase   :					Tuple        =    reversed_block_out_channels[i]

           _lowerCAmelCase   :					Optional[Any]        =    i == len(_a       ) - 1

           _lowerCAmelCase   :					Optional[int]        =    get_up_block(
               _a     ,						num_layers=self.layers_per_block + 1     ,						in_channels=_a     ,						out_channels=_a     ,						prev_output_channel=_a     ,						add_upsample=not is_final_block     ,						resnet_eps=1E-6     ,						resnet_act_fn=_a     ,						resnet_groups=_a     ,						attention_head_dim=_a     ,						temb_channels=_a     ,						resnet_time_scale_shift=_a     ,						)
           self.up_blocks.append(_a       )
           _lowerCAmelCase   :					str        =    output_channel

      # out
      if norm_type == "spatial":
           _lowerCAmelCase   :					str        =    SpatialNorm(block_out_channels[0]     ,						_a       )
      else:
           _lowerCAmelCase   :					int        =    nn.GroupNorm(num_channels=block_out_channels[0]     ,						num_groups=_a     ,						eps=1E-6       )
      _lowerCAmelCase   :					Dict        =    nn.SiLU()
      _lowerCAmelCase   :					str        =    nn.Convad(block_out_channels[0]     ,						_a     ,						3     ,						padding=1       )

      _lowerCAmelCase   :					int        =    False


 def    __magic_name__       (		self :	List[Any]     ,						A_ :	Any     ,						A_ :	List[Any]=None       ):






      '''simple docstring'''




      _lowerCAmelCase   :					Union[str, Any]        =    z
      _lowerCAmelCase   :					int        =    self.conv_in(_a       )

      _lowerCAmelCase   :					Any        =    next(iter(self.up_blocks.parameters()       )       ).dtype
      if self.training and self.gradient_checkpointing:

           def create_custom_forward(A_ :	Optional[Any]       ):
                def custom_forward(*A_ :	List[Any]       ):
                     return module(*_a       )

                return custom_forward

           if is_torch_version(">="     ,						"1.11.0"       ):
                # middle
                _lowerCAmelCase   :					List[str]        =    torch.utils.checkpoint.checkpoint(
                    create_custom_forward(self.mid_block       )     ,						_a     ,						_a     ,						use_reentrant=_a       )
                _lowerCAmelCase   :					List[Any]        =    sample.to(_a       )

                # up
                for up_block in self.up_blocks:
                     _lowerCAmelCase   :					Optional[int]        =    torch.utils.checkpoint.checkpoint(
                         create_custom_forward(_a       )     ,						_a     ,						_a     ,						use_reentrant=_a       )
           else:
                # middle
                _lowerCAmelCase   :					int        =    torch.utils.checkpoint.checkpoint(
                    create_custom_forward(self.mid_block       )     ,						_a     ,						_a       )
                _lowerCAmelCase   :					int        =    sample.to(_a       )

                # up
                for up_block in self.up_blocks:
                     _lowerCAmelCase   :					int        =    torch.utils.checkpoint.checkpoint(create_custom_forward(_a       )     ,						_a     ,						_a       )
      else:
           # middle
           _lowerCAmelCase   :					Any        =    self.mid_block(_a     ,						_a       )
           _lowerCAmelCase   :					Union[str, Any]        =    sample.to(_a       )

           # up
           for up_block in self.up_blocks:
                _lowerCAmelCase   :					Union[str, Any]        =    up_block(_a     ,						_a       )

        # post-process
      if latent_embeds is None:
           _lowerCAmelCase   :					Any        =    self.conv_norm_out(_a       )
      else:
           _lowerCAmelCase   :					Union[str, Any]        =    self.conv_norm_out(_a     ,						_a       )
      _lowerCAmelCase   :					str        =    self.conv_act(_a       )
      _lowerCAmelCase   :					Tuple        =    self.conv_out(_a       )

      return sample





class A__    ( nn.Module ):



 """simple docstring"""







 def __init__(		self :	str     ,						A_ :	str     ,						A_ :	Optional[Any]     ,						A_ :	List[str]     ,						A_ :	str=None     ,						A_ :	Union[str, Any]="random"     ,						A_ :	Any=False     ,						A_ :	Tuple=True       ):






      '''simple docstring'''




      super().__init__()
      _lowerCAmelCase   :					Tuple        =    n_e
      _lowerCAmelCase   :					Optional[Any]        =    vq_embed_dim
      _lowerCAmelCase   :					Tuple        =    beta
      _lowerCAmelCase   :					Dict        =    legacy

      _lowerCAmelCase   :					List[str]        =    nn.Embedding(self.n_e     ,						self.vq_embed_dim       )
      self.embedding.weight.data.uniform_(-1.0 / self.n_e     ,						1.0 / self.n_e       )

      _lowerCAmelCase   :					Tuple        =    remap
      if self.remap is not None:
           self.register_buffer("used"     ,						torch.tensor(np.load(self.remap       )       )       )
           _lowerCAmelCase   :					int        =    self.used.shape[0]
           _lowerCAmelCase   :					Optional[int]        =    unknown_index  # "random" or "extra" or integer
           if self.unknown_index == "extra":
                _lowerCAmelCase   :					Optional[int]        =    self.re_embed
                _lowerCAmelCase   :					Optional[Any]        =    self.re_embed + 1
           print(
               F"""Remapping {self.n_e} indices to {self.re_embed} indices. """
               F"""Using {self.unknown_index} for unknown indices."""       )
      else:
           _lowerCAmelCase   :					Dict        =    n_e

      _lowerCAmelCase   :					Dict        =    sane_index_shape

 def    __magic_name__       (		self :	Any     ,						A_ :	List[Any]       ):






      '''simple docstring'''




      _lowerCAmelCase   :					Tuple        =    inds.shape
      assert len(_a       ) > 1
      _lowerCAmelCase   :					Union[str, Any]        =    inds.reshape(ishape[0]     ,						-1       )
      _lowerCAmelCase   :					Union[str, Any]        =    self.used.to(_a       )
      _lowerCAmelCase   :					Optional[Any]        =    (inds[:, :, None] == used[None, None, ...]).long()
      _lowerCAmelCase   :					int        =    match.argmax(-1       )
      _lowerCAmelCase   :					Any        =    match.sum(2       ) < 1
      if self.unknown_index == "random":
           _lowerCAmelCase   :					List[str]        =    torch.randint(0     ,						self.re_embed     ,						size=new[unknown].shape       ).to(device=new.device       )
      else:
           _lowerCAmelCase   :					int        =    self.unknown_index
      return new.reshape(_a       )

 def    __magic_name__       (		self :	Any     ,						A_ :	Dict       ):






      '''simple docstring'''




      _lowerCAmelCase   :					Tuple        =    inds.shape
      assert len(_a       ) > 1
      _lowerCAmelCase   :					List[Any]        =    inds.reshape(ishape[0]     ,						-1       )
      _lowerCAmelCase   :					Union[str, Any]        =    self.used.to(_a       )
      if self.re_embed > self.used.shape[0]:  # extra token
           _lowerCAmelCase   :					str        =    0  # simply set to zero
      _lowerCAmelCase   :					Any        =    torch.gather(used[None, :][inds.shape[0] * [0], :]     ,						1     ,						_a       )
      return back.reshape(_a       )

 def    __magic_name__       (		self :	Union[str, Any]     ,						A_ :	Dict       ):






      '''simple docstring'''




      _lowerCAmelCase   :					Any        =    z.permute(0     ,						2     ,						3     ,						1       ).contiguous()
      _lowerCAmelCase   :					int        =    z.view(-1     ,						self.vq_embed_dim       )

      # distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z
      _lowerCAmelCase   :					Tuple        =    torch.argmin(torch.cdist(_a     ,						self.embedding.weight       )     ,						dim=1       )

      _lowerCAmelCase   :					Any        =    self.embedding(_a       ).view(z.shape       )
      _lowerCAmelCase   :					Optional[Any]        =    None
      _lowerCAmelCase   :					Any        =    None

      # compute loss for embedding
      if not self.legacy:
           _lowerCAmelCase   :					List[str]        =    self.beta * torch.mean((z_q.detach() - z) ** 2       ) + torch.mean((z_q - z.detach()) ** 2       )
      else:
           _lowerCAmelCase   :					Any        =    torch.mean((z_q.detach() - z) ** 2       ) + self.beta * torch.mean((z_q - z.detach()) ** 2       )

      # preserve gradients
      _lowerCAmelCase   :					List[str]        =    z + (z_q - z).detach()

      # reshape back to match original input shape
      _lowerCAmelCase   :					Dict        =    z_q.permute(0     ,						3     ,						1     ,						2       ).contiguous()

      if self.remap is not None:
           _lowerCAmelCase   :					str        =    min_encoding_indices.reshape(z.shape[0]     ,						-1       )  # add batch axis
           _lowerCAmelCase   :					int        =    self.remap_to_used(_a       )
           _lowerCAmelCase   :					List[Any]        =    min_encoding_indices.reshape(-1     ,						1       )  # flatten

      if self.sane_index_shape:
           _lowerCAmelCase   :					int        =    min_encoding_indices.reshape(z_q.shape[0]     ,						z_q.shape[2]     ,						z_q.shape[3]       )

      return z_q, loss, (perplexity, min_encodings, min_encoding_indices)


 def    __magic_name__       (		self :	int     ,						A_ :	Union[str, Any]     ,						A_ :	List[str]       ):






      '''simple docstring'''




      if self.remap is not None:
           _lowerCAmelCase   :					List[str]        =    indices.reshape(shape[0]     ,						-1       )  # add batch axis
           _lowerCAmelCase   :					Optional[int]        =    self.unmap_to_all(_a       )
           _lowerCAmelCase   :					Any        =    indices.reshape(-1       )  # flatten again

      # get quantized latent vectors
      _lowerCAmelCase   :					str        =    self.embedding(_a       )

      if shape is not None:
           _lowerCAmelCase   :					int        =    z_q.view(_a       )
           # reshape back to match original input shape
           _lowerCAmelCase   :					Optional[Any]        =    z_q.permute(0     ,						3     ,						1     ,						2       ).contiguous()

      return z_q





class A__    ( UpperCamelCase__ ):



 """simple docstring"""







 def __init__(		self :	Optional[Any]     ,						A_ :	str     ,						A_ :	int=False       ):






      '''simple docstring'''




      _lowerCAmelCase   :					Dict        =    parameters
      _lowerCAmelCase   :					Tuple        =    torch.chunk(_a     ,						2     ,						dim=1       )
      _lowerCAmelCase   :					Optional[Any]        =    torch.clamp(self.logvar     ,						-30.0     ,						20.0       )
      _lowerCAmelCase   :					Tuple        =    deterministic
      _lowerCAmelCase   :					Any        =    torch.exp(0.5 * self.logvar       )
      _lowerCAmelCase   :					Optional[int]        =    torch.exp(self.logvar       )
      if self.deterministic:
           _lowerCAmelCase   :					List[Any]        =    torch.zeros_like(
               self.mean     ,						device=self.parameters.device     ,						dtype=self.parameters.dtype       )

 def    __magic_name__       (		self :	Dict     ,						A_ :	Union[str, Any] = None       ):






      '''simple docstring'''




      _lowerCAmelCase   :					Union[str, Any]        =    randn_tensor(
          self.mean.shape     ,						generator=_a     ,						device=self.parameters.device     ,						dtype=self.parameters.dtype       )
      _lowerCAmelCase   :					int        =    self.mean + self.std * sample
      return x

 def    __magic_name__       (		self :	Optional[int]     ,						A_ :	int=None       ):






      '''simple docstring'''




      if self.deterministic:
           return torch.Tensor([0.0]       )
      else:
           if other is None:
                return 0.5 * torch.sum(torch.pow(self.mean     ,						2       ) + self.var - 1.0 - self.logvar     ,						dim=[1, 2, 3]       )
           else:
                return 0.5 * torch.sum(
                    torch.pow(self.mean - other.mean     ,						2       ) / other.var
                    + self.var / other.var
                    - 1.0
                    - self.logvar
                    + other.logvar     ,						dim=[1, 2, 3]     ,						)

 def    __magic_name__       (		self :	List[Any]     ,						A_ :	Optional[int]     ,						A_ :	Optional[Any]=[1, 2, 3]       ):






      '''simple docstring'''




      if self.deterministic:
           return torch.Tensor([0.0]       )
      _lowerCAmelCase   :					List[str]        =    np.log(2.0 * np.pi       )
      return 0.5 * torch.sum(logtwopi + self.logvar + torch.pow(sample - self.mean     ,						2       ) / self.var     ,						dim=_a       )


 def    __magic_name__       (		self :	str       ):






      '''simple docstring'''




      return self.mean









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


__UpperCAmelCase    				=     logging.get_logger(__name__)





class A__    ( A ):



 """simple docstring"""







 _lowercase     :					List[Any]      =						['''pixel_values''']

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






      '''simple docstring'''




      super().__init__(**A_       )
      _lowerCAmelCase   :					Optional[int]        =    size if size is not None else {"shortest_edge": 2_2_4}
      _lowerCAmelCase   :					Optional[int]        =    get_size_dict(A_     ,						default_to_square=A_       )
      _lowerCAmelCase   :					Any        =    crop_size if crop_size is not None else {"height": 2_5_6, "width": 2_5_6}
      _lowerCAmelCase   :					str        =    get_size_dict(A_     ,						param_name="crop_size"       )

      _lowerCAmelCase   :					Any        =    do_resize
      _lowerCAmelCase   :					Optional[Any]        =    size
      _lowerCAmelCase   :					str        =    resample
      _lowerCAmelCase   :					Optional[Any]        =    do_rescale
      _lowerCAmelCase   :					Dict        =    rescale_factor
      _lowerCAmelCase   :					Any        =    do_center_crop
      _lowerCAmelCase   :					List[Any]        =    crop_size
      _lowerCAmelCase   :					List[Any]        =    do_flip_channel_order

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






      '''simple docstring'''




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

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






      '''simple docstring'''




      _lowerCAmelCase   :					str        =    get_size_dict(A_       )
      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(A_     ,						size=(size["height"], size["width"])     ,						data_format=A_     ,						**A_       )

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






      '''simple docstring'''




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

 def    __magic_name__       (		self :	Optional[Any]     ,						A_ :	np.ndarray     ,						A_ :	Optional[Union[str, ChannelDimension]] = None       ):






      '''simple docstring'''




      return flip_channel_order(A_     ,						data_format=A_       )

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






      '''simple docstring'''




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

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

      _lowerCAmelCase   :					Dict        =    make_list_of_images(A_       )

      if not valid_images(A_       ):
           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.
      _lowerCAmelCase   :					Optional[int]        =    [to_numpy_array(A_       ) for image in images]

      if do_resize:
           _lowerCAmelCase   :					Any        =    [self.resize(image=A_     ,						size=A_     ,						resample=A_       ) for image in images]

      if do_center_crop:
           _lowerCAmelCase   :					Tuple        =    [self.center_crop(image=A_     ,						size=A_       ) for image in images]

      if do_rescale:
           _lowerCAmelCase   :					Any        =    [self.rescale(image=A_     ,						scale=A_       ) for image in images]

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

      _lowerCAmelCase   :					Optional[Any]        =    [to_channel_dimension_format(A_     ,						A_       ) for image in images]

      _lowerCAmelCase   :					Tuple        =    {"pixel_values": images}
      return BatchFeature(data=A_     ,						tensor_type=A_       )


 def    __magic_name__       (		self :	List[Any]     ,						A_ :	Union[str, Any]     ,						A_ :	List[Tuple] = None       ):






      '''simple docstring'''




      _lowerCAmelCase   :					Union[str, Any]        =    outputs.logits

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

           if is_torch_tensor(A_       ):
                _lowerCAmelCase   :					Dict        =    target_sizes.numpy()

           _lowerCAmelCase   :					List[Any]        =    []

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

      return semantic_segmentation






