





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 						_lowerCAmelCase(						_a):

  """simple docstring"""







  lowerCamelCase__			     =    42







class 						_lowerCAmelCase(						_a		,    _a):

  """simple docstring"""







  @register_to_config
  def __init__(					self  ,				UpperCAmelCase = 16  ,				UpperCAmelCase = 88  ,				UpperCAmelCase = None  ,				UpperCAmelCase = None  ,				UpperCAmelCase = 1  ,				UpperCAmelCase = 0.0  ,				UpperCAmelCase = 32  ,				UpperCAmelCase = None  ,				UpperCAmelCase = False  ,				UpperCAmelCase = None  ,				UpperCAmelCase = None  ,				UpperCAmelCase = None  ,				UpperCAmelCase = "geglu"  ,				UpperCAmelCase = None  ,				UpperCAmelCase = False  ,				UpperCAmelCase = False  ,				UpperCAmelCase = False  ,				UpperCAmelCase = "layer_norm"  ,				UpperCAmelCase = True  ,				)->					Optional[int]:
     super().__init__()
     __A      					=   use_linear_projection
     __A      					=   num_attention_heads
     __A      					=   attention_head_dim
     __A      					=   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
     __A      					=   (in_channels is not None) and (patch_size is None)
     __A      					=   num_vector_embeds is not None
     __A      					=   in_channels is not None and patch_size is not None

     if norm_type == "layer_norm" and num_embeds_ada_norm is not None:
        __A      					=   (
            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'''  ,				_lowercase  ,				standard_warn=_lowercase	)
        __A      					=   '''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:
        __A      					=   in_channels

        __A      					=   torch.nn.GroupNorm(num_groups=_lowercase  ,				num_channels=_lowercase  ,				eps=1e-6  ,				affine=_lowercase	)
        if use_linear_projection:
           __A      					=   nn.Linear(_lowercase  ,				_lowercase	)
        else:
           __A      					=   nn.Convad(_lowercase  ,				_lowercase  ,				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"

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

        __A      					=   ImagePositionalEmbeddings(
            num_embed=_lowercase  ,				embed_dim=_lowercase  ,				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"

        __A      					=   sample_size
        __A      					=   sample_size

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

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

     # 4. Define output layers
     __A      					=   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:
           __A      					=   nn.Linear(_lowercase  ,				_lowercase	)
        else:
           __A      					=   nn.Convad(_lowercase  ,				_lowercase  ,				kernel_size=1  ,				stride=1  ,				padding=0	)
     elif self.is_input_vectorized:
        __A      					=   nn.LayerNorm(_lowercase	)
        __A      					=   nn.Linear(_lowercase  ,				self.num_vector_embeds - 1	)
     elif self.is_input_patches:
        __A      					=   nn.LayerNorm(_lowercase  ,				elementwise_affine=_lowercase  ,				eps=1e-6	)
        __A      					=   nn.Linear(_lowercase  ,				2 * inner_dim	)
        __A      					=   nn.Linear(_lowercase  ,				patch_size * patch_size * self.out_channels	)





  def 		SCREAMING_SNAKE_CASE__       (					self  ,				UpperCAmelCase  ,				UpperCAmelCase = None  ,				UpperCAmelCase = None  ,				UpperCAmelCase = None  ,				UpperCAmelCase = None  ,				UpperCAmelCase = None  ,				UpperCAmelCase = None  ,				UpperCAmelCase = True  ,				)->					Optional[int]:
     # ensure attention_mask is a bias, and give it a singleton query_tokens dimension.
     #   we may have done this conversion already, e.g. if we came here via UNet2DConditionModel#forward.
     #   we can tell by counting dims; if ndim == 2: it's a mask rather than a bias.
     # expects mask of shape:
     #   [batch, key_tokens]
     # adds singleton query_tokens dimension:
     #   [batch,                    1, key_tokens]
     # this helps to broadcast it as a bias over attention scores, which will be in one of the following shapes:
     #   [batch,  heads, query_tokens, key_tokens] (e.g. torch sdp attn)
     #   [batch * heads, query_tokens, key_tokens] (e.g. xformers or classic attn)
     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)
        __A      					=   (1 - attention_mask.to(hidden_states.dtype	)) * -1_00_00.0
        __A      					=   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:
        __A      					=   (1 - encoder_attention_mask.to(hidden_states.dtype	)) * -1_00_00.0
        __A      					=   encoder_attention_mask.unsqueeze(1	)

     # 1. Input
     if self.is_input_continuous:
        __A					,						__A					,						__A					,						__A      					=   hidden_states.shape
        __A      					=   hidden_states

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

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

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

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

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

        # unpatchify
        __A      					=   __A      					=   int(hidden_states.shape[1] ** 0.5	)
        __A      					=   hidden_states.reshape(
            shape=(-1, height, width, self.patch_size, self.patch_size, self.out_channels)	)
        __A      					=   torch.einsum('''nhwpqc->nchpwq'''  ,				_lowercase	)
        __A      					=   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=_lowercase	)




from functools import partial
from typing import Any, Callable, Dict, List, Type, TypeVar, Union, overload

import torch
import torch.nn as nn
import torch.types




def        lowerCAmelCase(	SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       )->		torch.Tensor:



    """simple docstring"""






    if not inplace:
        UpperCamelCase_          =		ma + ma
    else:
        ma += ma

    return ma




def        lowerCAmelCase(	SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       )->		torch.Tensor:



    """simple docstring"""






    UpperCamelCase_          =		-1 * len(SCREAMING_SNAKE_CASE_       )
    UpperCamelCase_          =		list(range(len(tensor.shape[:zero_index]       )       )       )
    return tensor.permute(first_inds + [zero_index + i for i in inds]       )




def        lowerCAmelCase(	SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       )->		torch.Tensor:



    """simple docstring"""






    return t.reshape(t.shape[:-no_dims] + (-1,)       )




def        lowerCAmelCase(	SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_ = 1E-4       )->		torch.Tensor:



    """simple docstring"""






    UpperCamelCase_          =		mask.expand(*value.shape       )
    return torch.sum(mask * value       ,		dim=SCREAMING_SNAKE_CASE_       ) / (eps + torch.sum(SCREAMING_SNAKE_CASE_       ,		dim=SCREAMING_SNAKE_CASE_       ))




def        lowerCAmelCase(	SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_ = 2.3125       ,		SCREAMING_SNAKE_CASE_ = 21.6875       ,		SCREAMING_SNAKE_CASE_ = 6_4       )->		torch.Tensor:



    """simple docstring"""






    UpperCamelCase_          =		torch.linspace(SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       ,		no_bins - 1       ,		device=pts.device       )
    UpperCamelCase_          =		torch.sqrt(torch.sum((pts.unsqueeze(-2       ) - pts.unsqueeze(-3       )) ** 2       ,		dim=-1       )       )
    return torch.bucketize(SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       )




def        lowerCAmelCase(	SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       )->		dict:



    """simple docstring"""






    UpperCamelCase_          =		dicts[0]
    UpperCamelCase_          =		{}
    for k, v in first.items():
        UpperCamelCase_          =		[d[k] for d in dicts]
        if isinstance(SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       ):
            UpperCamelCase_          =		dict_multimap(SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       )
        else:
            UpperCamelCase_          =		fn(SCREAMING_SNAKE_CASE_       )

    return new_dict




def        lowerCAmelCase(	SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       )->		torch.Tensor:



    """simple docstring"""






    UpperCamelCase_          =		v_bins.view(((1,) * len(x.shape       )) + (len(SCREAMING_SNAKE_CASE_       ),)       )
    UpperCamelCase_          =		x[..., None] - reshaped_bins
    UpperCamelCase_          =		torch.argmin(torch.abs(SCREAMING_SNAKE_CASE_       )       ,		dim=-1       )
    return nn.functional.one_hot(SCREAMING_SNAKE_CASE_       ,		num_classes=len(SCREAMING_SNAKE_CASE_       )       ).float()




def        lowerCAmelCase(	SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_ = 0       ,		SCREAMING_SNAKE_CASE_ = 0       )->		torch.Tensor:



    """simple docstring"""






    UpperCamelCase_          =		[]
    for i, s in enumerate(data.shape[:no_batch_dims]       ):
        UpperCamelCase_          =		torch.arange(SCREAMING_SNAKE_CASE_       )
        UpperCamelCase_          =		r.view(*(*((1,) * i), -1, *((1,) * (len(inds.shape       ) - i - 1)))       )
        ranges.append(SCREAMING_SNAKE_CASE_       )

    UpperCamelCase_          =		[slice(SCREAMING_SNAKE_CASE_       ) for _ in range(len(data.shape       ) - no_batch_dims       )]
    UpperCamelCase_          =		inds
    ranges.extend(SCREAMING_SNAKE_CASE_       )
    # Matt note: Editing this to get around the behaviour of using a list as an array index changing
    # in recent Numpy versions
    return data[tuple(SCREAMING_SNAKE_CASE_       )]


SCREAMING_SNAKE_CASE						:Dict      =			TypeVar("""T""")




def        lowerCAmelCase(	SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       )->		Dict[Any, Union[dict, list, tuple, Any]]:



    """simple docstring"""






    UpperCamelCase_          =		{}
    for k, v in dic.items():
        if isinstance(SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       ):
            UpperCamelCase_          =		dict_map(SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       )
        else:
            UpperCamelCase_          =		tree_map(SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       )

    return new_dict




@overload
def        lowerCAmelCase(	SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       )->		Any:



    """simple docstring"""






    ...




@overload
def        lowerCAmelCase(	SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       )->		dict:



    """simple docstring"""






    ...




@overload
def        lowerCAmelCase(	SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       )->		list:



    """simple docstring"""






    ...




@overload
def        lowerCAmelCase(	SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       )->		tuple:



    """simple docstring"""






    ...




def        lowerCAmelCase(	SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       )->		Optional[int]:



    """simple docstring"""






    if isinstance(SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       ):
        return dict_map(SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       )
    elif isinstance(SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       ):
        return [tree_map(SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       ) for x in tree]
    elif isinstance(SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       ):
        return tuple(tree_map(SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       ) for x in tree       )
    elif isinstance(SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       ):
        return fn(SCREAMING_SNAKE_CASE_       )
    else:
        print(type(SCREAMING_SNAKE_CASE_       )       )
        raise ValueError("Not supported"       )
SCREAMING_SNAKE_CASE						:Dict      =			partial(tree_map, leaf_type=torch.Tensor)
