





'''simple docstring'''




from dataclasses import dataclass
from typing import Optional, Tuple, Union

import torch
import torch.nn as nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput
from .embeddings import GaussianFourierProjection, TimestepEmbedding, Timesteps
from .modeling_utils import ModelMixin
from .unet_ad_blocks import get_down_block, get_mid_block, get_out_block, get_up_block




@dataclass
class      UpperCamelCase__					(					lowercase_			):




					"""simple docstring"""
					SCREAMING_SNAKE_CASE__				=	42







class      UpperCamelCase__					(					lowercase_						,					lowercase_			):




					"""simple docstring"""


					@register_to_config
					def __init__(						self				:			int      ,						lowerCamelCase_				:			int = 6_55_36      ,						lowerCamelCase_				:			Optional[int] = None      ,						lowerCamelCase_				:			int = 2      ,						lowerCamelCase_				:			int = 2      ,						lowerCamelCase_				:			int = 0      ,						lowerCamelCase_				:			str = "fourier"      ,						lowerCamelCase_				:			bool = True      ,						lowerCamelCase_				:			bool = False      ,						lowerCamelCase_				:			float = 0.0      ,						lowerCamelCase_				:			Tuple[str] = ("DownBlock1DNoSkip", "DownBlock1D", "AttnDownBlock1D")      ,						lowerCamelCase_				:			Tuple[str] = ("AttnUpBlock1D", "UpBlock1D", "UpBlock1DNoSkip")      ,						lowerCamelCase_				:			Tuple[str] = "UNetMidBlock1D"      ,						lowerCamelCase_				:			str = None      ,						lowerCamelCase_				:			Tuple[int] = (32, 32, 64)      ,						lowerCamelCase_				:			str = None      ,						lowerCamelCase_				:			int = 8      ,						lowerCamelCase_				:			int = 1      ,						lowerCamelCase_				:			bool = False      ,						):

							'''simple docstring'''



							super().__init__()
							SCREAMING_SNAKE_CASE :				Optional[int]		=			sample_size

							# time
							if time_embedding_type == "fourier":
									SCREAMING_SNAKE_CASE :				Union[str, Any]		=			GaussianFourierProjection(
									    embedding_size=8      ,						set_W_to_weight=lowerCamelCase_      ,						log=lowerCamelCase_      ,						flip_sin_to_cos=lowerCamelCase_						)
									SCREAMING_SNAKE_CASE :				int		=			2 * block_out_channels[0]
							elif time_embedding_type == "positional":
									SCREAMING_SNAKE_CASE :				Union[str, Any]		=			Timesteps(
									    block_out_channels[0]      ,						flip_sin_to_cos=lowerCamelCase_      ,						downscale_freq_shift=lowerCamelCase_						)
									SCREAMING_SNAKE_CASE :				Optional[int]		=			block_out_channels[0]

							if use_timestep_embedding:
									SCREAMING_SNAKE_CASE :				Optional[int]		=			block_out_channels[0] * 4
									SCREAMING_SNAKE_CASE :				Optional[int]		=			TimestepEmbedding(
									    in_channels=lowerCamelCase_      ,						time_embed_dim=lowerCamelCase_      ,						act_fn=lowerCamelCase_      ,						out_dim=block_out_channels[0]      ,						)

							SCREAMING_SNAKE_CASE :				Optional[int]		=			nn.ModuleList([]						)
							SCREAMING_SNAKE_CASE :				List[str]		=			None
							SCREAMING_SNAKE_CASE :				Optional[int]		=			nn.ModuleList([]						)
							SCREAMING_SNAKE_CASE :				Tuple		=			None

							# down
							SCREAMING_SNAKE_CASE :				Optional[int]		=			in_channels
							for i, down_block_type in enumerate(lowerCamelCase_						):
									SCREAMING_SNAKE_CASE :				Optional[Any]		=			output_channel
									SCREAMING_SNAKE_CASE :				int		=			block_out_channels[i]

									if i == 0:
											input_channel += extra_in_channels

									SCREAMING_SNAKE_CASE :				Tuple		=			i == len(lowerCamelCase_						) - 1

									SCREAMING_SNAKE_CASE :				List[str]		=			get_down_block(
									    lowerCamelCase_      ,						num_layers=lowerCamelCase_      ,						in_channels=lowerCamelCase_      ,						out_channels=lowerCamelCase_      ,						temb_channels=block_out_channels[0]      ,						add_downsample=not is_final_block or downsample_each_block      ,						)
									self.down_blocks.append(lowerCamelCase_						)

							# mid
							SCREAMING_SNAKE_CASE :				Tuple		=			get_mid_block(
							    lowerCamelCase_      ,						in_channels=block_out_channels[-1]      ,						mid_channels=block_out_channels[-1]      ,						out_channels=block_out_channels[-1]      ,						embed_dim=block_out_channels[0]      ,						num_layers=lowerCamelCase_      ,						add_downsample=lowerCamelCase_      ,						)

							# up
							SCREAMING_SNAKE_CASE :				int		=			list(reversed(lowerCamelCase_						)						)
							SCREAMING_SNAKE_CASE :				Optional[Any]		=			reversed_block_out_channels[0]
							if out_block_type is None:
									SCREAMING_SNAKE_CASE :				Tuple		=			out_channels
							else:
									SCREAMING_SNAKE_CASE :				List[str]		=			block_out_channels[0]

							for i, up_block_type in enumerate(lowerCamelCase_						):
									SCREAMING_SNAKE_CASE :				int		=			output_channel
									SCREAMING_SNAKE_CASE :				Optional[Any]		=			(
									    reversed_block_out_channels[i + 1] if i < len(lowerCamelCase_						) - 1 else final_upsample_channels
									)

									SCREAMING_SNAKE_CASE :				Dict		=			i == len(lowerCamelCase_						) - 1

									SCREAMING_SNAKE_CASE :				Any		=			get_up_block(
									    lowerCamelCase_      ,						num_layers=lowerCamelCase_      ,						in_channels=lowerCamelCase_      ,						out_channels=lowerCamelCase_      ,						temb_channels=block_out_channels[0]      ,						add_upsample=not is_final_block      ,						)
									self.up_blocks.append(lowerCamelCase_						)
									SCREAMING_SNAKE_CASE :				Optional[int]		=			output_channel

							# out
							SCREAMING_SNAKE_CASE :				Optional[int]		=			norm_num_groups if norm_num_groups is not None else min(block_out_channels[0] // 4      ,						32						)
							SCREAMING_SNAKE_CASE :				List[str]		=			get_out_block(
							    out_block_type=lowerCamelCase_      ,						num_groups_out=lowerCamelCase_      ,						embed_dim=block_out_channels[0]      ,						out_channels=lowerCamelCase_      ,						act_fn=lowerCamelCase_      ,						fc_dim=block_out_channels[-1] // 4      ,						)



					def lowerCamelCase_			(						self				:			Dict      ,						lowerCamelCase_				:			torch.FloatTensor      ,						lowerCamelCase_				:			Union[torch.Tensor, float, int]      ,						lowerCamelCase_				:			bool = True      ,						):

							'''simple docstring'''



							SCREAMING_SNAKE_CASE :				Optional[Any]		=			timestep
							if not torch.is_tensor(lowerCamelCase_						):
									SCREAMING_SNAKE_CASE :				Dict		=			torch.tensor([timesteps]      ,						dtype=torch.long      ,						device=sample.device						)
							elif torch.is_tensor(lowerCamelCase_						) and len(timesteps.shape						) == 0:
									SCREAMING_SNAKE_CASE :				List[str]		=			timesteps[None].to(sample.device						)

							SCREAMING_SNAKE_CASE :				List[str]		=			self.time_proj(lowerCamelCase_						)
							if self.config.use_timestep_embedding:
									SCREAMING_SNAKE_CASE :				List[Any]		=			self.time_mlp(lowerCamelCase_						)
							else:
									SCREAMING_SNAKE_CASE :				List[Any]		=			timestep_embed[..., None]
									SCREAMING_SNAKE_CASE :				int		=			timestep_embed.repeat([1, 1, sample.shape[2]]						).to(sample.dtype						)
									SCREAMING_SNAKE_CASE :				Union[str, Any]		=			timestep_embed.broadcast_to((sample.shape[:1] + timestep_embed.shape[1:])						)

							# 2. down
							SCREAMING_SNAKE_CASE :				int		=			()
							for downsample_block in self.down_blocks:
									SCREAMING_SNAKE_CASE,							SCREAMING_SNAKE_CASE :				Any		=			downsample_block(hidden_states=lowerCamelCase_      ,						temb=lowerCamelCase_						)
									down_block_res_samples += res_samples

							# 3. mid
							if self.mid_block:
									SCREAMING_SNAKE_CASE :				Optional[int]		=			self.mid_block(lowerCamelCase_      ,						lowerCamelCase_						)

							# 4. up
							for i, upsample_block in enumerate(self.up_blocks						):
									SCREAMING_SNAKE_CASE :				Union[str, Any]		=			down_block_res_samples[-1:]
									SCREAMING_SNAKE_CASE :				Any		=			down_block_res_samples[:-1]
									SCREAMING_SNAKE_CASE :				Dict		=			upsample_block(lowerCamelCase_      ,						res_hidden_states_tuple=lowerCamelCase_      ,						temb=lowerCamelCase_						)

							# 5. post-process
							if self.out_block:
									SCREAMING_SNAKE_CASE :				Any		=			self.out_block(lowerCamelCase_      ,						lowerCamelCase_						)

							if not return_dict:
									return (sample,)

							return UNetaDOutput(sample=lowerCamelCase_						)










'''simple docstring'''




import unittest

from transformers import DonutProcessor


__UpperCAmelCase            =       """naver-clova-ix/donut-base"""




class      UpperCamelCase__					(					unittest.TestCase			):




					"""simple docstring"""


					def lowerCamelCase_			(						self				:			List[str]						):

							'''simple docstring'''



							SCREAMING_SNAKE_CASE :				Any		=			DonutProcessor.from_pretrained(lowerCamelCase_						)



					def lowerCamelCase_			(						self				:			Tuple						):

							'''simple docstring'''



							SCREAMING_SNAKE_CASE :				Optional[int]		=			{
							    """name""": """John Doe""",
							    """age""": """99""",
							    """city""": """Atlanta""",
							    """state""": """GA""",
							    """zip""": """30301""",
							    """phone""": """123-4567""",
							    """nicknames""": [{"""nickname""": """Johnny"""}, {"""nickname""": """JD"""}],
							}

							SCREAMING_SNAKE_CASE :				List[Any]		=			(
							    """<s_name>John Doe</s_name><s_age>99</s_age><s_city>Atlanta</s_city>"""
							    """<s_state>GA</s_state><s_zip>30301</s_zip><s_phone>123-4567</s_phone>"""
							    """<s_nicknames><s_nickname>Johnny</s_nickname>"""
							    """<sep/><s_nickname>JD</s_nickname></s_nicknames>"""
							)
							SCREAMING_SNAKE_CASE :				List[Any]		=			self.processor.tokenajson(lowerCamelCase_						)

							self.assertDictEqual(lowerCamelCase_      ,						lowerCamelCase_						)



