
"""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      UpperCamelCase	(       __magic_name__      ):

				"""simple docstring"""



				UpperCAmelCase_				:      torch.FloatTensor






class      UpperCamelCase	(       __magic_name__  ,      __magic_name__      ):

				"""simple docstring"""



				@register_to_config
				def __init__(				self		,	lowercase__ = 16		,	lowercase__ = 88		,	lowercase__ = None		,	lowercase__ = None		,	lowercase__ = 1		,	lowercase__ = 0.0		,	lowercase__ = 32		,	lowercase__ = None		,	lowercase__ = False		,	lowercase__ = None		,	lowercase__ = None		,	lowercase__ = None		,	lowercase__ = "geglu"		,	lowercase__ = None		,	lowercase__ = False		,	lowercase__ = False		,	lowercase__ = False		,	lowercase__ = "layer_norm"		,	lowercase__ = True		,	)    ->    Dict:
									"""simple docstring"""







									super().__init__()
									SCREAMING_SNAKE_CASE				=	use_linear_projection
									SCREAMING_SNAKE_CASE				=	num_attention_heads
									SCREAMING_SNAKE_CASE				=	attention_head_dim
									SCREAMING_SNAKE_CASE				=	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
									SCREAMING_SNAKE_CASE				=	(in_channels is not None) and (patch_size is None)
									SCREAMING_SNAKE_CASE				=	num_vector_embeds is not None
									SCREAMING_SNAKE_CASE				=	in_channels is not None and patch_size is not None

									if norm_type == "layer_norm" and num_embeds_ada_norm is not None:
														SCREAMING_SNAKE_CASE				=	(
														    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__			)
														SCREAMING_SNAKE_CASE				=	'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:
														SCREAMING_SNAKE_CASE				=	in_channels

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

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

														SCREAMING_SNAKE_CASE				=	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"

														SCREAMING_SNAKE_CASE				=	sample_size
														SCREAMING_SNAKE_CASE				=	sample_size

														SCREAMING_SNAKE_CASE				=	patch_size
														SCREAMING_SNAKE_CASE				=	PatchEmbed(
														    height=lowercase__		,	width=lowercase__		,	patch_size=lowercase__		,	in_channels=lowercase__		,	embed_dim=lowercase__		,	)

									# 3. Define transformers blocks
									SCREAMING_SNAKE_CASE				=	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
									SCREAMING_SNAKE_CASE				=	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:
																			SCREAMING_SNAKE_CASE				=	nn.Linear(lowercase__		,	lowercase__			)
														else:
																			SCREAMING_SNAKE_CASE				=	nn.Convad(lowercase__		,	lowercase__		,	kernel_size=1		,	stride=1		,	padding=0			)
									elif self.is_input_vectorized:
														SCREAMING_SNAKE_CASE				=	nn.LayerNorm(lowercase__			)
														SCREAMING_SNAKE_CASE				=	nn.Linear(lowercase__		,	self.num_vector_embeds - 1			)
									elif self.is_input_patches:
														SCREAMING_SNAKE_CASE				=	nn.LayerNorm(lowercase__		,	elementwise_affine=lowercase__		,	eps=1E-6			)
														SCREAMING_SNAKE_CASE				=	nn.Linear(lowercase__		,	2 * inner_dim			)
														SCREAMING_SNAKE_CASE				=	nn.Linear(lowercase__		,	patch_size * patch_size * self.out_channels			)




				def 							A					(				self		,	lowercase__		,	lowercase__ = None		,	lowercase__ = None		,	lowercase__ = None		,	lowercase__ = None		,	lowercase__ = None		,	lowercase__ = None		,	lowercase__ = True		,	)    ->    List[Any]:
									"""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)
														SCREAMING_SNAKE_CASE				=	(1 - attention_mask.to(hidden_states.dtype			)) * -10_000.0
														SCREAMING_SNAKE_CASE				=	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:
														SCREAMING_SNAKE_CASE				=	(1 - encoder_attention_mask.to(hidden_states.dtype			)) * -10_000.0
														SCREAMING_SNAKE_CASE				=	encoder_attention_mask.unsqueeze(1			)

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

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

									# 2. Blocks
									for block in self.transformer_blocks:
														SCREAMING_SNAKE_CASE				=	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:
																			SCREAMING_SNAKE_CASE				=	hidden_states.reshape(lowercase__		,	lowercase__		,	lowercase__		,	lowercase__			).permute(0		,	3		,	1		,	2			).contiguous()
																			SCREAMING_SNAKE_CASE				=	self.proj_out(lowercase__			)
														else:
																			SCREAMING_SNAKE_CASE				=	self.proj_out(lowercase__			)
																			SCREAMING_SNAKE_CASE				=	hidden_states.reshape(lowercase__		,	lowercase__		,	lowercase__		,	lowercase__			).permute(0		,	3		,	1		,	2			).contiguous()

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

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

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


'''simple docstring'''



import textwrap

import pyarrow as pa
import pytest

from datasets import Features, Image
from datasets.packaged_modules.text.text import Text

from ..utils import require_pil
@pytest.fixture
def       lowerCAmelCase_		( SCREAMING_SNAKE_CASE_  )   ->   Union[str, Any]:




   """simple docstring"""
   _SCREAMING_SNAKE_CASE          = tmp_path / """text.txt"""
   _SCREAMING_SNAKE_CASE          = textwrap.dedent(
       """\
        Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua.
        Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat.
        Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur.
        Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum.

        Second paragraph:
        Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua.
        Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat.
        Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur.
        Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum.
        """  )
   with open(SCREAMING_SNAKE_CASE_			,					"""w"""			,					encoding="""utf-8"""  ) as f:
      f.write(SCREAMING_SNAKE_CASE_  )
   return str(SCREAMING_SNAKE_CASE_  )
@pytest.fixture
def       lowerCAmelCase_		( SCREAMING_SNAKE_CASE_			,					SCREAMING_SNAKE_CASE_  )   ->   Optional[int]:




   """simple docstring"""
   _SCREAMING_SNAKE_CASE          = tmp_path / """text_with_image.txt"""
   with open(SCREAMING_SNAKE_CASE_			,					"""w"""			,					encoding="""utf-8"""  ) as f:
      f.write(SCREAMING_SNAKE_CASE_  )
   return str(SCREAMING_SNAKE_CASE_  )
@pytest.mark.parametrize("""keep_linebreaks"""			,					[True, False]  )
def       lowerCAmelCase_		( SCREAMING_SNAKE_CASE_			,					SCREAMING_SNAKE_CASE_  )   ->   int:




   """simple docstring"""
   with open(SCREAMING_SNAKE_CASE_			,					encoding="""utf-8"""  ) as f:
      _SCREAMING_SNAKE_CASE          = f.read().splitlines(keepends=SCREAMING_SNAKE_CASE_  )
   _SCREAMING_SNAKE_CASE          = Text(keep_linebreaks=SCREAMING_SNAKE_CASE_			,					encoding="""utf-8"""  )
   _SCREAMING_SNAKE_CASE          = text._generate_tables([[text_file]]  )
   _SCREAMING_SNAKE_CASE          = pa.concat_tables([table for _, table in generator]  ).to_pydict()["""text"""]
   assert generated_content == expected_content
@require_pil
def       lowerCAmelCase_		( SCREAMING_SNAKE_CASE_  )   ->   List[str]:




   """simple docstring"""
   with open(SCREAMING_SNAKE_CASE_			,					encoding="""utf-8"""  ) as f:
      _SCREAMING_SNAKE_CASE          = f.read().splitlines()[0]
   _SCREAMING_SNAKE_CASE          = Text(encoding="""utf-8"""			,					features=Features({"""image""": Image()}  )  )
   _SCREAMING_SNAKE_CASE          = text._generate_tables([[text_file_with_image]]  )
   _SCREAMING_SNAKE_CASE          = pa.concat_tables([table for _, table in generator]  )
   assert pa_table.schema.field("""image"""  ).type == Image()()
   _SCREAMING_SNAKE_CASE          = pa_table.to_pydict()["""image"""]
   assert generated_content == [{"path": image_file, "bytes": None}]







@pytest.mark.parametrize("""sample_by"""			,					["""line""", """paragraph""", """document"""]  )
def       lowerCAmelCase_		( SCREAMING_SNAKE_CASE_			,					SCREAMING_SNAKE_CASE_  )   ->   List[Any]:




   """simple docstring"""
   with open(SCREAMING_SNAKE_CASE_			,					encoding="""utf-8"""  ) as f:
      _SCREAMING_SNAKE_CASE          = f.read()
   if sample_by == "line":
      _SCREAMING_SNAKE_CASE          = expected_content.splitlines()
   elif sample_by == "paragraph":
      _SCREAMING_SNAKE_CASE          = expected_content.split("""\n\n"""  )
   elif sample_by == "document":
      _SCREAMING_SNAKE_CASE          = [expected_content]
   _SCREAMING_SNAKE_CASE          = Text(sample_by=SCREAMING_SNAKE_CASE_			,					encoding="""utf-8"""			,					chunksize=1_00  )
   _SCREAMING_SNAKE_CASE          = text._generate_tables([[text_file]]  )
   _SCREAMING_SNAKE_CASE          = pa.concat_tables([table for _, table in generator]  ).to_pydict()["""text"""]
   assert generated_content == expected_content






