
import functools
import math

import flax.linen as nn
import jax
import jax.numpy as jnp
def        __lowercase(							__snake_case				:		str    ,__snake_case				:		List[Any]    ,__snake_case				:		Union[str, Any]    ,__snake_case				:		List[Any]    ,__snake_case				:		int = 40_96					)       ->   Tuple:
       __snake_case       ,       __snake_case       ,       __snake_case						       =   key.shape[-3:]
       __snake_case						       =   value.shape[-1]
       __snake_case						       =   min(__snake_case    ,__snake_case					)
       __snake_case						       =   query / jnp.sqrt(__snake_case					)

       @functools.partial(jax.checkpoint    ,prevent_cse=__snake_case					)
       def summarize_chunk(__snake_case				:		Dict    ,__snake_case				:		Optional[Any]    ,__snake_case				:		int					):
              __snake_case						       =   jnp.einsum('...qhd,...khd->...qhk'    ,__snake_case    ,__snake_case    ,precision=__snake_case					)

              __snake_case						       =   jnp.max(__snake_case    ,axis=-1    ,keepdims=__snake_case					)
              __snake_case						       =   jax.lax.stop_gradient(__snake_case					)
              __snake_case						       =   jnp.exp(attn_weights - max_score					)

              __snake_case						       =   jnp.einsum('...vhf,...qhv->...qhf'    ,__snake_case    ,__snake_case    ,precision=__snake_case					)
              __snake_case						       =   jnp.einsum('...qhk->...qh'    ,__snake_case					)

              return (exp_values, exp_weights.sum(axis=-1					), max_score)

       def chunk_scanner(__snake_case				:		int					):
              # julienne key array
              __snake_case						       =   jax.lax.dynamic_slice(
                  operand=__snake_case    ,start_indices=[0] * (key.ndim - 3) + [chunk_idx, 0, 0]    ,slice_sizes=list(key.shape[:-3]					) + [key_chunk_size, num_heads, k_features]    ,)

              # julienne value array
              __snake_case						       =   jax.lax.dynamic_slice(
                  operand=__snake_case    ,start_indices=[0] * (value.ndim - 3) + [chunk_idx, 0, 0]    ,slice_sizes=list(value.shape[:-3]					) + [key_chunk_size, num_heads, v_features]    ,)

              return summarize_chunk(__snake_case    ,__snake_case    ,__snake_case					)

       __snake_case       ,       __snake_case       ,       __snake_case						       =   jax.lax.map(f=__snake_case    ,xs=jnp.arange(0    ,__snake_case    ,__snake_case					)					)

       __snake_case						       =   jnp.max(__snake_case    ,axis=0    ,keepdims=__snake_case					)
       __snake_case						       =   jnp.exp(chunk_max - global_max					)

       chunk_values *= jnp.expand_dims(__snake_case    ,axis=-1					)
       chunk_weights *= max_diffs

       __snake_case						       =   chunk_values.sum(axis=0					)
       __snake_case						       =   jnp.expand_dims(__snake_case    ,-1					).sum(axis=0					)

       return all_values / all_weights
def        __lowercase(							__snake_case				:		Dict    ,__snake_case				:		Optional[Any]    ,__snake_case				:		Tuple    ,__snake_case				:		List[str]=jax.lax.Precision.HIGHEST    ,__snake_case				:		int = 10_24    ,__snake_case				:		int = 40_96					)       ->   str:
       __snake_case       ,       __snake_case       ,       __snake_case						       =   query.shape[-3:]

       def chunk_scanner(__snake_case				:		Tuple    ,__snake_case				:		str					):
              # julienne query array
              __snake_case						       =   jax.lax.dynamic_slice(
                  operand=__snake_case    ,start_indices=([0] * (query.ndim - 3)) + [chunk_idx, 0, 0]    ,slice_sizes=list(query.shape[:-3]					) + [min(__snake_case    ,__snake_case					), num_heads, q_features]    ,)

              return (
                  chunk_idx + query_chunk_size,  # unused ignore it
                  _query_chunk_attention(
                      query=__snake_case    ,key=__snake_case    ,value=__snake_case    ,precision=__snake_case    ,key_chunk_size=__snake_case					),
              )

       __snake_case       ,       __snake_case						       =   jax.lax.scan(
           f=__snake_case    ,init=0    ,xs=__snake_case    ,length=math.ceil(num_q / query_chunk_size					)					)

       return jnp.concatenate(__snake_case    ,axis=-3					)  # fuse the chunked result back







class  _lowerCamelCase			(nn.Module					):
     lowercase__				  =			42
     lowercase__				  =			8
     lowercase__				  =			64
     lowercase__				  =			0.0
     lowercase__				  =			False
     lowercase__				  =			jnp.floataa

     def 						__lowerCamelCase				(  self    ):
            __snake_case						       =   self.dim_head * self.heads
            __snake_case						       =   self.dim_head**-0.5

            # Weights were exported with old names {to_q, to_k, to_v, to_out}
            __snake_case						       =   nn.Dense(SCREAMING_SNAKE_CASE_							,       use_bias=SCREAMING_SNAKE_CASE_							,       dtype=self.dtype							,       name='to_q'    )
            __snake_case						       =   nn.Dense(SCREAMING_SNAKE_CASE_							,       use_bias=SCREAMING_SNAKE_CASE_							,       dtype=self.dtype							,       name='to_k'    )
            __snake_case						       =   nn.Dense(SCREAMING_SNAKE_CASE_							,       use_bias=SCREAMING_SNAKE_CASE_							,       dtype=self.dtype							,       name='to_v'    )

            __snake_case						       =   nn.Dense(self.query_dim							,       dtype=self.dtype							,       name='to_out_0'    )
            __snake_case						       =   nn.Dropout(rate=self.dropout    )

     def 						__lowerCamelCase				(  self							,       SCREAMING_SNAKE_CASE_    ):
            __snake_case       ,       __snake_case       ,       __snake_case						       =   tensor.shape
            __snake_case						       =   self.heads
            __snake_case						       =   tensor.reshape(SCREAMING_SNAKE_CASE_							,       SCREAMING_SNAKE_CASE_							,       SCREAMING_SNAKE_CASE_							,       dim // head_size    )
            __snake_case						       =   jnp.transpose(SCREAMING_SNAKE_CASE_							,       (0, 2, 1, 3)    )
            __snake_case						       =   tensor.reshape(batch_size * head_size							,       SCREAMING_SNAKE_CASE_							,       dim // head_size    )
            return tensor

     def 						__lowerCamelCase				(  self							,       SCREAMING_SNAKE_CASE_    ):
            __snake_case       ,       __snake_case       ,       __snake_case						       =   tensor.shape
            __snake_case						       =   self.heads
            __snake_case						       =   tensor.reshape(batch_size // head_size							,       SCREAMING_SNAKE_CASE_							,       SCREAMING_SNAKE_CASE_							,       SCREAMING_SNAKE_CASE_    )
            __snake_case						       =   jnp.transpose(SCREAMING_SNAKE_CASE_							,       (0, 2, 1, 3)    )
            __snake_case						       =   tensor.reshape(batch_size // head_size							,       SCREAMING_SNAKE_CASE_							,       dim * head_size    )
            return tensor

     def __call__(  self							,       SCREAMING_SNAKE_CASE_							,       SCREAMING_SNAKE_CASE_=None							,       SCREAMING_SNAKE_CASE_=True    ):
            __snake_case						       =   hidden_states if context is None else context

            __snake_case						       =   self.query(SCREAMING_SNAKE_CASE_    )
            __snake_case						       =   self.key(SCREAMING_SNAKE_CASE_    )
            __snake_case						       =   self.value(SCREAMING_SNAKE_CASE_    )

            __snake_case						       =   self.reshape_heads_to_batch_dim(SCREAMING_SNAKE_CASE_    )
            __snake_case						       =   self.reshape_heads_to_batch_dim(SCREAMING_SNAKE_CASE_    )
            __snake_case						       =   self.reshape_heads_to_batch_dim(SCREAMING_SNAKE_CASE_    )

            if self.use_memory_efficient_attention:
                   __snake_case						       =   query_states.transpose(1							,       0							,       2    )
                   __snake_case						       =   key_states.transpose(1							,       0							,       2    )
                   __snake_case						       =   value_states.transpose(1							,       0							,       2    )

                   # this if statement create a chunk size for each layer of the unet
                   # the chunk size is equal to the query_length dimension of the deepest layer of the unet

                   __snake_case						       =   query_states.shape[-3]
                   if flatten_latent_dim % 64 == 0:
                          __snake_case						       =   int(flatten_latent_dim / 64    )
                   elif flatten_latent_dim % 16 == 0:
                          __snake_case						       =   int(flatten_latent_dim / 16    )
                   elif flatten_latent_dim % 4 == 0:
                          __snake_case						       =   int(flatten_latent_dim / 4    )
                   else:
                          __snake_case						       =   int(SCREAMING_SNAKE_CASE_    )

                   __snake_case						       =   jax_memory_efficient_attention(
                       SCREAMING_SNAKE_CASE_							,       SCREAMING_SNAKE_CASE_							,       SCREAMING_SNAKE_CASE_							,       query_chunk_size=SCREAMING_SNAKE_CASE_							,       key_chunk_size=4_096 * 4    )

                   __snake_case						       =   hidden_states.transpose(1							,       0							,       2    )
            else:
                   # compute attentions
                   __snake_case						       =   jnp.einsum('b i d, b j d->b i j'							,       SCREAMING_SNAKE_CASE_							,       SCREAMING_SNAKE_CASE_    )
                   __snake_case						       =   attention_scores * self.scale
                   __snake_case						       =   nn.softmax(SCREAMING_SNAKE_CASE_							,       axis=2    )

                   # attend to values
                   __snake_case						       =   jnp.einsum('b i j, b j d -> b i d'							,       SCREAMING_SNAKE_CASE_							,       SCREAMING_SNAKE_CASE_    )

            __snake_case						       =   self.reshape_batch_dim_to_heads(SCREAMING_SNAKE_CASE_    )
            __snake_case						       =   self.proj_attn(SCREAMING_SNAKE_CASE_    )
            return self.dropout_layer(SCREAMING_SNAKE_CASE_							,       deterministic=SCREAMING_SNAKE_CASE_    )







class  _lowerCamelCase			(nn.Module					):
     lowercase__				  =			42
     lowercase__				  =			42
     lowercase__				  =			42
     lowercase__				  =			0.0
     lowercase__				  =			False
     lowercase__				  =			jnp.floataa
     lowercase__				  =			False

     def 						__lowerCamelCase				(  self    ):
            # self attention (or cross_attention if only_cross_attention is True)
            __snake_case						       =   FlaxAttention(
                self.dim							,       self.n_heads							,       self.d_head							,       self.dropout							,       self.use_memory_efficient_attention							,       dtype=self.dtype    )
            # cross attention
            __snake_case						       =   FlaxAttention(
                self.dim							,       self.n_heads							,       self.d_head							,       self.dropout							,       self.use_memory_efficient_attention							,       dtype=self.dtype    )
            __snake_case						       =   FlaxFeedForward(dim=self.dim							,       dropout=self.dropout							,       dtype=self.dtype    )
            __snake_case						       =   nn.LayerNorm(epsilon=1E-5							,       dtype=self.dtype    )
            __snake_case						       =   nn.LayerNorm(epsilon=1E-5							,       dtype=self.dtype    )
            __snake_case						       =   nn.LayerNorm(epsilon=1E-5							,       dtype=self.dtype    )
            __snake_case						       =   nn.Dropout(rate=self.dropout    )

     def __call__(  self							,       SCREAMING_SNAKE_CASE_							,       SCREAMING_SNAKE_CASE_							,       SCREAMING_SNAKE_CASE_=True    ):
            # self attention
            __snake_case						       =   hidden_states
            if self.only_cross_attention:
                   __snake_case						       =   self.attna(self.norma(SCREAMING_SNAKE_CASE_    )							,       SCREAMING_SNAKE_CASE_							,       deterministic=SCREAMING_SNAKE_CASE_    )
            else:
                   __snake_case						       =   self.attna(self.norma(SCREAMING_SNAKE_CASE_    )							,       deterministic=SCREAMING_SNAKE_CASE_    )
            __snake_case						       =   hidden_states + residual

            # cross attention
            __snake_case						       =   hidden_states
            __snake_case						       =   self.attna(self.norma(SCREAMING_SNAKE_CASE_    )							,       SCREAMING_SNAKE_CASE_							,       deterministic=SCREAMING_SNAKE_CASE_    )
            __snake_case						       =   hidden_states + residual

            # feed forward
            __snake_case						       =   hidden_states
            __snake_case						       =   self.ff(self.norma(SCREAMING_SNAKE_CASE_    )							,       deterministic=SCREAMING_SNAKE_CASE_    )
            __snake_case						       =   hidden_states + residual

            return self.dropout_layer(SCREAMING_SNAKE_CASE_							,       deterministic=SCREAMING_SNAKE_CASE_    )







class  _lowerCamelCase			(nn.Module					):
     lowercase__				  =			42
     lowercase__				  =			42
     lowercase__				  =			42
     lowercase__				  =			1
     lowercase__				  =			0.0
     lowercase__				  =			False
     lowercase__				  =			False
     lowercase__				  =			jnp.floataa
     lowercase__				  =			False

     def 						__lowerCamelCase				(  self    ):
            __snake_case						       =   nn.GroupNorm(num_groups=32							,       epsilon=1E-5    )

            __snake_case						       =   self.n_heads * self.d_head
            if self.use_linear_projection:
                   __snake_case						       =   nn.Dense(SCREAMING_SNAKE_CASE_							,       dtype=self.dtype    )
            else:
                   __snake_case						       =   nn.Conv(
                       SCREAMING_SNAKE_CASE_							,       kernel_size=(1, 1)							,       strides=(1, 1)							,       padding='VALID'							,       dtype=self.dtype							,       )

            __snake_case						       =   [
                FlaxBasicTransformerBlock(
                    SCREAMING_SNAKE_CASE_							,       self.n_heads							,       self.d_head							,       dropout=self.dropout							,       only_cross_attention=self.only_cross_attention							,       dtype=self.dtype							,       use_memory_efficient_attention=self.use_memory_efficient_attention							,       )
                for _ in range(self.depth    )
            ]

            if self.use_linear_projection:
                   __snake_case						       =   nn.Dense(SCREAMING_SNAKE_CASE_							,       dtype=self.dtype    )
            else:
                   __snake_case						       =   nn.Conv(
                       SCREAMING_SNAKE_CASE_							,       kernel_size=(1, 1)							,       strides=(1, 1)							,       padding='VALID'							,       dtype=self.dtype							,       )

            __snake_case						       =   nn.Dropout(rate=self.dropout    )

     def __call__(  self							,       SCREAMING_SNAKE_CASE_							,       SCREAMING_SNAKE_CASE_							,       SCREAMING_SNAKE_CASE_=True    ):
            __snake_case       ,       __snake_case       ,       __snake_case       ,       __snake_case						       =   hidden_states.shape
            __snake_case						       =   hidden_states
            __snake_case						       =   self.norm(SCREAMING_SNAKE_CASE_    )
            if self.use_linear_projection:
                   __snake_case						       =   hidden_states.reshape(SCREAMING_SNAKE_CASE_							,       height * width							,       SCREAMING_SNAKE_CASE_    )
                   __snake_case						       =   self.proj_in(SCREAMING_SNAKE_CASE_    )
            else:
                   __snake_case						       =   self.proj_in(SCREAMING_SNAKE_CASE_    )
                   __snake_case						       =   hidden_states.reshape(SCREAMING_SNAKE_CASE_							,       height * width							,       SCREAMING_SNAKE_CASE_    )

            for transformer_block in self.transformer_blocks:
                   __snake_case						       =   transformer_block(SCREAMING_SNAKE_CASE_							,       SCREAMING_SNAKE_CASE_							,       deterministic=SCREAMING_SNAKE_CASE_    )

            if self.use_linear_projection:
                   __snake_case						       =   self.proj_out(SCREAMING_SNAKE_CASE_    )
                   __snake_case						       =   hidden_states.reshape(SCREAMING_SNAKE_CASE_							,       SCREAMING_SNAKE_CASE_							,       SCREAMING_SNAKE_CASE_							,       SCREAMING_SNAKE_CASE_    )
            else:
                   __snake_case						       =   hidden_states.reshape(SCREAMING_SNAKE_CASE_							,       SCREAMING_SNAKE_CASE_							,       SCREAMING_SNAKE_CASE_							,       SCREAMING_SNAKE_CASE_    )
                   __snake_case						       =   self.proj_out(SCREAMING_SNAKE_CASE_    )

            __snake_case						       =   hidden_states + residual
            return self.dropout_layer(SCREAMING_SNAKE_CASE_							,       deterministic=SCREAMING_SNAKE_CASE_    )







class  _lowerCamelCase			(nn.Module					):
     lowercase__				  =			42
     lowercase__				  =			0.0
     lowercase__				  =			jnp.floataa

     def 						__lowerCamelCase				(  self    ):
            # The second linear layer needs to be called
            # net_2 for now to match the index of the Sequential layer
            __snake_case						       =   FlaxGEGLU(self.dim							,       self.dropout							,       self.dtype    )
            __snake_case						       =   nn.Dense(self.dim							,       dtype=self.dtype    )

     def __call__(  self							,       SCREAMING_SNAKE_CASE_							,       SCREAMING_SNAKE_CASE_=True    ):
            __snake_case						       =   self.net_a(SCREAMING_SNAKE_CASE_							,       deterministic=SCREAMING_SNAKE_CASE_    )
            __snake_case						       =   self.net_a(SCREAMING_SNAKE_CASE_    )
            return hidden_states







class  _lowerCamelCase			(nn.Module					):
     lowercase__				  =			42
     lowercase__				  =			0.0
     lowercase__				  =			jnp.floataa

     def 						__lowerCamelCase				(  self    ):
            __snake_case						       =   self.dim * 4
            __snake_case						       =   nn.Dense(inner_dim * 2							,       dtype=self.dtype    )
            __snake_case						       =   nn.Dropout(rate=self.dropout    )

     def __call__(  self							,       SCREAMING_SNAKE_CASE_							,       SCREAMING_SNAKE_CASE_=True    ):
            __snake_case						       =   self.proj(SCREAMING_SNAKE_CASE_    )
            __snake_case       ,       __snake_case						       =   jnp.split(SCREAMING_SNAKE_CASE_							,       2							,       axis=2    )
            return self.dropout_layer(hidden_linear * nn.gelu(SCREAMING_SNAKE_CASE_    )							,       deterministic=SCREAMING_SNAKE_CASE_    )


from typing import Dict, List, Optional, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import center_crop, normalize, rescale, resize, to_channel_dimension_format
from ...image_utils import (
    IMAGENET_STANDARD_MEAN,
    IMAGENET_STANDARD_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_vision_available, logging


if is_vision_available():
 import PIL


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







class  _lowerCamelCase			(lowerCamelCase					):
     lowercase__				  =			["""pixel_values"""]

     def __init__(  self							,       SCREAMING_SNAKE_CASE_ = True							,       SCREAMING_SNAKE_CASE_ = None							,       SCREAMING_SNAKE_CASE_ = PIL.Image.NEAREST							,       SCREAMING_SNAKE_CASE_ = False							,       SCREAMING_SNAKE_CASE_ = None							,       SCREAMING_SNAKE_CASE_ = 1 / 255							,       SCREAMING_SNAKE_CASE_ = False							,       SCREAMING_SNAKE_CASE_ = True							,       SCREAMING_SNAKE_CASE_ = True							,       SCREAMING_SNAKE_CASE_ = None							,       SCREAMING_SNAKE_CASE_ = None							,       SCREAMING_SNAKE_CASE_ = True							,       **SCREAMING_SNAKE_CASE_							,       ):
            super().__init__(**SCREAMING_SNAKE_CASE_    )
            __snake_case						       =   size if size is not None else {'height': 346, 'width': 346}
            __snake_case						       =   get_size_dict(SCREAMING_SNAKE_CASE_    )
            __snake_case						       =   crop_size if crop_size is not None else {'height': 289, 'width': 289}
            __snake_case						       =   get_size_dict(SCREAMING_SNAKE_CASE_							,       param_name='crop_size'    )

            __snake_case						       =   do_resize
            __snake_case						       =   size
            __snake_case						       =   resample
            __snake_case						       =   do_center_crop
            __snake_case						       =   crop_size
            __snake_case						       =   do_rescale
            __snake_case						       =   rescale_factor
            __snake_case						       =   rescale_offset
            __snake_case						       =   do_normalize
            __snake_case						       =   image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN
            __snake_case						       =   image_std if image_std is not None else IMAGENET_STANDARD_STD
            __snake_case						       =   include_top

     def 						__lowerCamelCase				(  self							,       SCREAMING_SNAKE_CASE_							,       SCREAMING_SNAKE_CASE_							,       SCREAMING_SNAKE_CASE_ = PIL.Image.NEAREST							,       SCREAMING_SNAKE_CASE_ = None							,       **SCREAMING_SNAKE_CASE_							,       ):
            __snake_case						       =   get_size_dict(SCREAMING_SNAKE_CASE_    )
            if "height" not in size or "width" not in size:
                   raise ValueError(f'''The size dictionary must have keys \'height\' and \'width\'. Got {size.keys()}'''    )
            return resize(
                SCREAMING_SNAKE_CASE_							,       size=(size['height'], size['width'])							,       resample=SCREAMING_SNAKE_CASE_							,       data_format=SCREAMING_SNAKE_CASE_							,       **SCREAMING_SNAKE_CASE_    )

     def 						__lowerCamelCase				(  self							,       SCREAMING_SNAKE_CASE_							,       SCREAMING_SNAKE_CASE_							,       SCREAMING_SNAKE_CASE_ = None							,       **SCREAMING_SNAKE_CASE_							,       ):
            __snake_case						       =   get_size_dict(SCREAMING_SNAKE_CASE_    )
            if "height" not in size or "width" not in size:
                   raise ValueError(f'''The size dictionary must have 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 						__lowerCamelCase				(  self							,       SCREAMING_SNAKE_CASE_							,       SCREAMING_SNAKE_CASE_							,       SCREAMING_SNAKE_CASE_ = True							,       SCREAMING_SNAKE_CASE_ = None							,       **SCREAMING_SNAKE_CASE_							,       ):
            if offset:
                   __snake_case						       =   (image - 1_2_7.5) * scale
                   if data_format is not None:
                          __snake_case						       =   to_channel_dimension_format(SCREAMING_SNAKE_CASE_							,       SCREAMING_SNAKE_CASE_    )
                   __snake_case						       =   rescaled_image.astype(np.floataa    )
            else:
                   __snake_case						       =   rescale(SCREAMING_SNAKE_CASE_							,       scale=SCREAMING_SNAKE_CASE_							,       data_format=SCREAMING_SNAKE_CASE_							,       **SCREAMING_SNAKE_CASE_    )
            return rescaled_image

     def 						__lowerCamelCase				(  self							,       SCREAMING_SNAKE_CASE_							,       SCREAMING_SNAKE_CASE_							,       SCREAMING_SNAKE_CASE_							,       SCREAMING_SNAKE_CASE_ = None							,       **SCREAMING_SNAKE_CASE_							,       ):
            return normalize(SCREAMING_SNAKE_CASE_							,       mean=SCREAMING_SNAKE_CASE_							,       std=SCREAMING_SNAKE_CASE_							,       data_format=SCREAMING_SNAKE_CASE_							,       **SCREAMING_SNAKE_CASE_    )

     def 						__lowerCamelCase				(  self							,       SCREAMING_SNAKE_CASE_							,       SCREAMING_SNAKE_CASE_ = None							,       SCREAMING_SNAKE_CASE_ = None							,       SCREAMING_SNAKE_CASE_=None							,       SCREAMING_SNAKE_CASE_ = None							,       SCREAMING_SNAKE_CASE_ = None							,       SCREAMING_SNAKE_CASE_ = None							,       SCREAMING_SNAKE_CASE_ = None							,       SCREAMING_SNAKE_CASE_ = None							,       SCREAMING_SNAKE_CASE_ = None							,       SCREAMING_SNAKE_CASE_ = None							,       SCREAMING_SNAKE_CASE_ = None							,       SCREAMING_SNAKE_CASE_ = None							,       SCREAMING_SNAKE_CASE_ = None							,       SCREAMING_SNAKE_CASE_ = ChannelDimension.FIRST							,       **SCREAMING_SNAKE_CASE_							,       ):
            __snake_case						       =   do_resize if do_resize is not None else self.do_resize
            __snake_case						       =   resample if resample is not None else self.resample
            __snake_case						       =   do_center_crop if do_center_crop is not None else self.do_center_crop
            __snake_case						       =   do_rescale if do_rescale is not None else self.do_rescale
            __snake_case						       =   rescale_factor if rescale_factor is not None else self.rescale_factor
            __snake_case						       =   rescale_offset if rescale_offset is not None else self.rescale_offset
            __snake_case						       =   do_normalize if do_normalize is not None else self.do_normalize
            __snake_case						       =   image_mean if image_mean is not None else self.image_mean
            __snake_case						       =   image_std if image_std is not None else self.image_std
            __snake_case						       =   include_top if include_top is not None else self.include_top

            __snake_case						       =   size if size is not None else self.size
            __snake_case						       =   get_size_dict(SCREAMING_SNAKE_CASE_    )
            __snake_case						       =   crop_size if crop_size is not None else self.crop_size
            __snake_case						       =   get_size_dict(SCREAMING_SNAKE_CASE_							,       param_name='crop_size'    )

            __snake_case						       =   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 or resample is None:
                   raise ValueError('Size and resample must be specified if do_resize is True.'    )

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

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

            if do_normalize and (image_mean is None or image_std is None):
                   raise ValueError('Image mean and std must be specified if do_normalize is True.'    )

            # All transformations expect numpy arrays.
            __snake_case						       =   [to_numpy_array(SCREAMING_SNAKE_CASE_    ) for image in images]

            if do_resize:
                   __snake_case						       =   [self.resize(image=SCREAMING_SNAKE_CASE_							,       size=SCREAMING_SNAKE_CASE_							,       resample=SCREAMING_SNAKE_CASE_    ) for image in images]

            if do_center_crop:
                   __snake_case						       =   [self.center_crop(image=SCREAMING_SNAKE_CASE_							,       size=SCREAMING_SNAKE_CASE_    ) for image in images]

            if do_rescale:
                   __snake_case						       =   [self.rescale(image=SCREAMING_SNAKE_CASE_							,       scale=SCREAMING_SNAKE_CASE_							,       offset=SCREAMING_SNAKE_CASE_    ) for image in images]

            if do_normalize:
                   __snake_case						       =   [self.normalize(image=SCREAMING_SNAKE_CASE_							,       mean=SCREAMING_SNAKE_CASE_							,       std=SCREAMING_SNAKE_CASE_    ) for image in images]

            if include_top:
                   __snake_case						       =   [self.normalize(image=SCREAMING_SNAKE_CASE_							,       mean=[0, 0, 0]							,       std=SCREAMING_SNAKE_CASE_    ) for image in images]

            __snake_case						       =   [to_channel_dimension_format(SCREAMING_SNAKE_CASE_							,       SCREAMING_SNAKE_CASE_    ) for image in images]

            __snake_case						       =   {'pixel_values': images}
            return BatchFeature(data=SCREAMING_SNAKE_CASE_							,       tensor_type=SCREAMING_SNAKE_CASE_    )
