






'''simple docstring'''





from dataclasses import dataclass
from typing import Any, Dict, Optional

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

from ..configuration_utils import ConfigMixin, register_to_config
from ..models.embeddings import ImagePositionalEmbeddings
from ..utils import BaseOutput, deprecate
from .attention import BasicTransformerBlock
from .embeddings import PatchEmbed
from .modeling_utils import ModelMixin

@dataclass
class 					snake_case_							(       UpperCamelCase__      ):







   """simple docstring"""



   __lowerCAmelCase      =						42






class 					snake_case_							(       UpperCamelCase__      ,			UpperCamelCase__      ):







   """simple docstring"""



   @register_to_config
   def __init__(    self:	List[str]   ,       snake_case__:	Dict = 16   ,       snake_case__:	Optional[int] = 88   ,       snake_case__:	Any = None   ,       snake_case__:	List[Any] = None   ,       snake_case__:	int = 1   ,       snake_case__:	Optional[int] = 0.0   ,       snake_case__:	Optional[int] = 32   ,       snake_case__:	Dict = None   ,       snake_case__:	Union[str, Any] = False   ,       snake_case__:	Tuple = None   ,       snake_case__:	Optional[Any] = None   ,       snake_case__:	Tuple = None   ,       snake_case__:	Any = "geglu"   ,       snake_case__:	Tuple = None   ,       snake_case__:	Tuple = False   ,       snake_case__:	Any = False   ,       snake_case__:	Optional[Any] = False   ,       snake_case__:	str = "layer_norm"   ,       snake_case__:	Optional[Any] = True   ,       )  ->		Tuple:


       """simple docstring"""



       super().__init__()
       _lowercase     =	use_linear_projection
       _lowercase     =	num_attention_heads
       _lowercase     =	attention_head_dim
       _lowercase     =	num_attention_heads * attention_head_dim

       # 1. Transformer2DModel can process both standard continuous images of shape `(batch_size, num_channels, width, height)` as well as quantized image embeddings of shape `(batch_size, num_image_vectors)`
       # Define whether input is continuous or discrete depending on configuration
       _lowercase     =	(in_channels is not None) and (patch_size is None)
       _lowercase     =	num_vector_embeds is not None
       _lowercase     =	in_channels is not None and patch_size is not None

       if norm_type == "layer_norm" and num_embeds_ada_norm is not None:
           _lowercase     =	(
               f'''The configuration file of this model: {self.__class__} is outdated. `norm_type` is either not set or'''
               ''' incorrectly set to `\'layer_norm\'`.Make sure to set `norm_type` to `\'ada_norm\'` in the config.'''
               ''' Please make sure to update the config accordingly as leaving `norm_type` might led to incorrect'''
               ''' results in future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it'''
               ''' would be very nice if you could open a Pull request for the `transformer/config.json` file'''
           )
           deprecate('norm_type!=num_embeds_ada_norm'   ,       '1.0.0'   ,       snake_case_   ,       standard_warn=snake_case_	)
           _lowercase     =	'''ada_norm'''

       if self.is_input_continuous and self.is_input_vectorized:
           raise ValueError(
               f'''Cannot define both `in_channels`: {in_channels} and `num_vector_embeds`: {num_vector_embeds}. Make'''
               ' sure that either `in_channels` or `num_vector_embeds` is None.'	)
       elif self.is_input_vectorized and self.is_input_patches:
           raise ValueError(
               f'''Cannot define both `num_vector_embeds`: {num_vector_embeds} and `patch_size`: {patch_size}. Make'''
               ' sure that either `num_vector_embeds` or `num_patches` is None.'	)
       elif not self.is_input_continuous and not self.is_input_vectorized and not self.is_input_patches:
           raise ValueError(
               f'''Has to define `in_channels`: {in_channels}, `num_vector_embeds`: {num_vector_embeds}, or patch_size:'''
               f''' {patch_size}. Make sure that `in_channels`, `num_vector_embeds` or `num_patches` is not None.'''	)

       # 2. Define input layers
       if self.is_input_continuous:
           _lowercase     =	in_channels

           _lowercase     =	torch.nn.GroupNorm(num_groups=snake_case_   ,       num_channels=snake_case_   ,       eps=1E-6   ,       affine=snake_case_	)
           if use_linear_projection:
               _lowercase     =	nn.Linear(snake_case_   ,       snake_case_	)
           else:
               _lowercase     =	nn.Convad(snake_case_   ,       snake_case_   ,       kernel_size=1   ,       stride=1   ,       padding=0	)
       elif self.is_input_vectorized:
           assert sample_size is not None, "Transformer2DModel over discrete input must provide sample_size"
           assert num_vector_embeds is not None, "Transformer2DModel over discrete input must provide num_embed"

           _lowercase     =	sample_size
           _lowercase     =	sample_size
           _lowercase     =	num_vector_embeds
           _lowercase     =	self.height * self.width

           _lowercase     =	ImagePositionalEmbeddings(
               num_embed=snake_case_   ,       embed_dim=snake_case_   ,       height=self.height   ,       width=self.width	)
       elif self.is_input_patches:
           assert sample_size is not None, "Transformer2DModel over patched input must provide sample_size"

           _lowercase     =	sample_size
           _lowercase     =	sample_size

           _lowercase     =	patch_size
           _lowercase     =	PatchEmbed(
               height=snake_case_   ,       width=snake_case_   ,       patch_size=snake_case_   ,       in_channels=snake_case_   ,       embed_dim=snake_case_   ,       )

       # 3. Define transformers blocks
       _lowercase     =	nn.ModuleList(
           [
               BasicTransformerBlock(
                   snake_case_   ,       snake_case_   ,       snake_case_   ,       dropout=snake_case_   ,       cross_attention_dim=snake_case_   ,       activation_fn=snake_case_   ,       num_embeds_ada_norm=snake_case_   ,       attention_bias=snake_case_   ,       only_cross_attention=snake_case_   ,       upcast_attention=snake_case_   ,       norm_type=snake_case_   ,       norm_elementwise_affine=snake_case_   ,       )
               for d in range(snake_case_	)
           ]	)

       # 4. Define output layers
       _lowercase     =	in_channels if out_channels is None else out_channels
       if self.is_input_continuous:
           # TODO: should use out_channels for continuous projections
           if use_linear_projection:
               _lowercase     =	nn.Linear(snake_case_   ,       snake_case_	)
           else:
               _lowercase     =	nn.Convad(snake_case_   ,       snake_case_   ,       kernel_size=1   ,       stride=1   ,       padding=0	)
       elif self.is_input_vectorized:
           _lowercase     =	nn.LayerNorm(snake_case_	)
           _lowercase     =	nn.Linear(snake_case_   ,       self.num_vector_embeds - 1	)
       elif self.is_input_patches:
           _lowercase     =	nn.LayerNorm(snake_case_   ,       elementwise_affine=snake_case_   ,       eps=1E-6	)
           _lowercase     =	nn.Linear(snake_case_   ,       2 * inner_dim	)
           _lowercase     =	nn.Linear(snake_case_   ,       patch_size * patch_size * self.out_channels	)

   def 			SCREAMING_SNAKE_CASE__				(    self:	int   ,       snake_case__:	str   ,       snake_case__:	Union[str, Any] = None   ,       snake_case__:	List[str] = None   ,       snake_case__:	List[str] = None   ,       snake_case__:	List[Any] = None   ,       snake_case__:	str = None   ,       snake_case__:	Optional[int] = None   ,       snake_case__:	str = True   ,       )  ->		Tuple:


       """simple docstring"""



       if attention_mask is not None and attention_mask.ndim == 2:
           # assume that mask is expressed as:
           #   (1 = keep,      0 = discard)
           # convert mask into a bias that can be added to attention scores:
           #       (keep = +0,     discard = -10000.0)
           _lowercase     =	(1 - attention_mask.to(hidden_states.dtype	)) * -10_000.0
           _lowercase     =	attention_mask.unsqueeze(1	)

       # convert encoder_attention_mask to a bias the same way we do for attention_mask
       if encoder_attention_mask is not None and encoder_attention_mask.ndim == 2:
           _lowercase     =	(1 - encoder_attention_mask.to(hidden_states.dtype	)) * -10_000.0
           _lowercase     =	encoder_attention_mask.unsqueeze(1	)

       # 1. Input
       if self.is_input_continuous:
           _lowercase     =	hidden_states.shape
           _lowercase     =	hidden_states

           _lowercase     =	self.norm(snake_case_	)
           if not self.use_linear_projection:
               _lowercase     =	self.proj_in(snake_case_	)
               _lowercase     =	hidden_states.shape[1]
               _lowercase     =	hidden_states.permute(0   ,       2   ,       3   ,       1	).reshape(snake_case_   ,       height * width   ,       snake_case_	)
           else:
               _lowercase     =	hidden_states.shape[1]
               _lowercase     =	hidden_states.permute(0   ,       2   ,       3   ,       1	).reshape(snake_case_   ,       height * width   ,       snake_case_	)
               _lowercase     =	self.proj_in(snake_case_	)
       elif self.is_input_vectorized:
           _lowercase     =	self.latent_image_embedding(snake_case_	)
       elif self.is_input_patches:
           _lowercase     =	self.pos_embed(snake_case_	)

       # 2. Blocks
       for block in self.transformer_blocks:
           _lowercase     =	block(
               snake_case_   ,       attention_mask=snake_case_   ,       encoder_hidden_states=snake_case_   ,       encoder_attention_mask=snake_case_   ,       timestep=snake_case_   ,       cross_attention_kwargs=snake_case_   ,       class_labels=snake_case_   ,       )

       # 3. Output
       if self.is_input_continuous:
           if not self.use_linear_projection:
               _lowercase     =	hidden_states.reshape(snake_case_   ,       snake_case_   ,       snake_case_   ,       snake_case_	).permute(0   ,       3   ,       1   ,       2	).contiguous()
               _lowercase     =	self.proj_out(snake_case_	)
           else:
               _lowercase     =	self.proj_out(snake_case_	)
               _lowercase     =	hidden_states.reshape(snake_case_   ,       snake_case_   ,       snake_case_   ,       snake_case_	).permute(0   ,       3   ,       1   ,       2	).contiguous()

           _lowercase     =	hidden_states + residual
       elif self.is_input_vectorized:
           _lowercase     =	self.norm_out(snake_case_	)
           _lowercase     =	self.out(snake_case_	)
           # (batch, self.num_vector_embeds - 1, self.num_latent_pixels)
           _lowercase     =	logits.permute(0   ,       2   ,       1	)

           # log(p(x_0))
           _lowercase     =	F.log_softmax(logits.double()   ,       dim=1	).float()
       elif self.is_input_patches:
           # TODO: cleanup!
           _lowercase     =	self.transformer_blocks[0].norma.emb(
               snake_case_   ,       snake_case_   ,       hidden_dtype=hidden_states.dtype	)
           _lowercase     =	self.proj_out_a(F.silu(snake_case_	)	).chunk(2   ,       dim=1	)
           _lowercase     =	self.norm_out(snake_case_	) * (1 + scale[:, None]) + shift[:, None]
           _lowercase     =	self.proj_out_a(snake_case_	)

           # unpatchify
           _lowercase     =	int(hidden_states.shape[1] ** 0.5	)
           _lowercase     =	hidden_states.reshape(
               shape=(-1, height, width, self.patch_size, self.patch_size, self.out_channels)	)
           _lowercase     =	torch.einsum('nhwpqc->nchpwq'   ,       snake_case_	)
           _lowercase     =	hidden_states.reshape(
               shape=(-1, self.out_channels, height * self.patch_size, width * self.patch_size)	)

       if not return_dict:
           return (output,)

       return TransformeraDModelOutput(sample=snake_case_	)








from __future__ import annotations
def        UpperCAmelCase      (   _lowercase,      _lowercase      ):


    '''simple docstring'''



    _lowercase     =	[None] * (len(_lowercase      ) + len(_lowercase      ))

    _lowercase     =	0  # pointer to current index for left Half
    _lowercase     =	0  # pointer to current index for the right Half
    _lowercase     =	0  # pointer to current index for the sorted array Half

    while pointera < len(_lowercase      ) and pointera < len(_lowercase      ):
        if left_half[pointera] < right_half[pointera]:
            _lowercase     =	left_half[pointera]
            pointera += 1
            index += 1
        else:
            _lowercase     =	right_half[pointera]
            pointera += 1
            index += 1
    while pointera < len(_lowercase      ):
        _lowercase     =	left_half[pointera]
        pointera += 1
        index += 1

    while pointera < len(_lowercase      ):
        _lowercase     =	right_half[pointera]
        pointera += 1
        index += 1

    return sorted_array
def        UpperCAmelCase      (   _lowercase      ):


    '''simple docstring'''



    if len(_lowercase      ) <= 1:
        return array
    # the actual formula to calculate the middle element = left + (right - left) // 2
    # this avoids integer overflow in case of large N
    _lowercase     =	0 + (len(_lowercase      ) - 0) // 2

    # Split the array into halves till the array length becomes equal to One
    # merge the arrays of single length returned by mergeSort function and
    # pass them into the merge arrays function which merges the array
    _lowercase     =	array[:middle]
    _lowercase     =	array[middle:]

    return merge(merge_sort(_lowercase      ),      merge_sort(_lowercase      )      )


if __name__ == "__main__":
   import doctest

   doctest.testmod()
