




from typing import TYPE_CHECKING

from ...utils import _LazyModule


_lowerCamelCase    				=		{'tokenization_bert_japanese': ['BertJapaneseTokenizer', 'CharacterTokenizer', 'MecabTokenizer']}


if TYPE_CHECKING:
 from .tokenization_bert_japanese import BertJapaneseTokenizer, CharacterTokenizer, MecabTokenizer

else:
 import sys

 _lowerCamelCase    				=		_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)











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 UNetMidBlockaD, get_down_block, get_up_block



@dataclass
class     a  (     _A    ):

     '''simple docstring'''





     lowerCAmelCase       : torch.FloatTensor



class     a  (     _A			,     _A    ):

     '''simple docstring'''





     @register_to_config
     def __init__(			self     :    Optional[Any]      ,					__snake_case     :    Optional[Union[int, Tuple[int, int]]] = None      ,					__snake_case     :    int = 3      ,					__snake_case     :    int = 3      ,					__snake_case     :    bool = False      ,					__snake_case     :    str = "positional"      ,					__snake_case     :    int = 0      ,					__snake_case     :    bool = True      ,					__snake_case     :    Tuple[str] = ("DownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D")      ,					__snake_case     :    Tuple[str] = ("AttnUpBlock2D", "AttnUpBlock2D", "AttnUpBlock2D", "UpBlock2D")      ,					__snake_case     :    Tuple[int] = (2_24, 4_48, 6_72, 8_96)      ,					__snake_case     :    int = 2      ,					__snake_case     :    float = 1      ,					__snake_case     :    int = 1      ,					__snake_case     :    str = "conv"      ,					__snake_case     :    str = "conv"      ,					__snake_case     :    str = "silu"      ,					__snake_case     :    Optional[int] = 8      ,					__snake_case     :    int = 32      ,					__snake_case     :    float = 1E-5      ,					__snake_case     :    str = "default"      ,					__snake_case     :    bool = True      ,					__snake_case     :    Optional[str] = None      ,					__snake_case     :    Optional[int] = None      ,					):
         super().__init__()

         UpperCAmelCase_											=  sample_size
         UpperCAmelCase_											=  block_out_channels[0] * 4

         # Check inputs
         if len(__snake_case       ) != len(__snake_case       ):
             raise ValueError(
                 F'Must provide the same number of `down_block_types` as `up_block_types`. `down_block_types`: {down_block_types}. `up_block_types`: {up_block_types}.'       )

         if len(__snake_case       ) != len(__snake_case       ):
             raise ValueError(
                 F'Must provide the same number of `block_out_channels` as `down_block_types`. `block_out_channels`: {block_out_channels}. `down_block_types`: {down_block_types}.'       )

         # input
         UpperCAmelCase_											=  nn.Convad(__snake_case      ,					block_out_channels[0]      ,					kernel_size=3      ,					padding=(1, 1)       )

         # time
         if time_embedding_type == "fourier":
             UpperCAmelCase_											=  GaussianFourierProjection(embedding_size=block_out_channels[0]      ,					scale=16       )
             UpperCAmelCase_											=  2 * block_out_channels[0]
         elif time_embedding_type == "positional":
             UpperCAmelCase_											=  Timesteps(block_out_channels[0]      ,					__snake_case      ,					__snake_case       )
             UpperCAmelCase_											=  block_out_channels[0]

         UpperCAmelCase_											=  TimestepEmbedding(__snake_case      ,					__snake_case       )

         # class embedding
         if class_embed_type is None and num_class_embeds is not None:
             UpperCAmelCase_											=  nn.Embedding(__snake_case      ,					__snake_case       )
         elif class_embed_type == "timestep":
             UpperCAmelCase_											=  TimestepEmbedding(__snake_case      ,					__snake_case       )
         elif class_embed_type == "identity":
             UpperCAmelCase_											=  nn.Identity(__snake_case      ,					__snake_case       )
         else:
             UpperCAmelCase_											=  None

         UpperCAmelCase_											=  nn.ModuleList([]       )
         UpperCAmelCase_											=  None
         UpperCAmelCase_											=  nn.ModuleList([]       )

         # down
         UpperCAmelCase_											=  block_out_channels[0]
         for i, down_block_type in enumerate(__snake_case       ):
             UpperCAmelCase_											=  output_channel
             UpperCAmelCase_											=  block_out_channels[i]
             UpperCAmelCase_											=  i == len(__snake_case       ) - 1

             UpperCAmelCase_											=  get_down_block(
                 __snake_case      ,					num_layers=__snake_case      ,					in_channels=__snake_case      ,					out_channels=__snake_case      ,					temb_channels=__snake_case      ,					add_downsample=not is_final_block      ,					resnet_eps=__snake_case      ,					resnet_act_fn=__snake_case      ,					resnet_groups=__snake_case      ,					attention_head_dim=attention_head_dim if attention_head_dim is not None else output_channel      ,					downsample_padding=__snake_case      ,					resnet_time_scale_shift=__snake_case      ,					downsample_type=__snake_case      ,					)
             self.down_blocks.append(__snake_case       )

         # mid
         UpperCAmelCase_											=  UNetMidBlockaD(
             in_channels=block_out_channels[-1]      ,					temb_channels=__snake_case      ,					resnet_eps=__snake_case      ,					resnet_act_fn=__snake_case      ,					output_scale_factor=__snake_case      ,					resnet_time_scale_shift=__snake_case      ,					attention_head_dim=attention_head_dim if attention_head_dim is not None else block_out_channels[-1]      ,					resnet_groups=__snake_case      ,					add_attention=__snake_case      ,					)

         # up
         UpperCAmelCase_											=  list(reversed(__snake_case       )       )
         UpperCAmelCase_											=  reversed_block_out_channels[0]
         for i, up_block_type in enumerate(__snake_case       ):
             UpperCAmelCase_											=  output_channel
             UpperCAmelCase_											=  reversed_block_out_channels[i]
             UpperCAmelCase_											=  reversed_block_out_channels[min(i + 1      ,					len(__snake_case       ) - 1       )]

             UpperCAmelCase_											=  i == len(__snake_case       ) - 1

             UpperCAmelCase_											=  get_up_block(
                 __snake_case      ,					num_layers=layers_per_block + 1      ,					in_channels=__snake_case      ,					out_channels=__snake_case      ,					prev_output_channel=__snake_case      ,					temb_channels=__snake_case      ,					add_upsample=not is_final_block      ,					resnet_eps=__snake_case      ,					resnet_act_fn=__snake_case      ,					resnet_groups=__snake_case      ,					attention_head_dim=attention_head_dim if attention_head_dim is not None else output_channel      ,					resnet_time_scale_shift=__snake_case      ,					upsample_type=__snake_case      ,					)
             self.up_blocks.append(__snake_case       )
             UpperCAmelCase_											=  output_channel

         # out
         UpperCAmelCase_											=  norm_num_groups if norm_num_groups is not None else min(block_out_channels[0] // 4      ,					32       )
         UpperCAmelCase_											=  nn.GroupNorm(num_channels=block_out_channels[0]      ,					num_groups=__snake_case      ,					eps=__snake_case       )
         UpperCAmelCase_											=  nn.SiLU()
         UpperCAmelCase_											=  nn.Convad(block_out_channels[0]      ,					__snake_case      ,					kernel_size=3      ,					padding=1       )




     def 						lowerCamelCase_     (			self     :    Dict      ,					__snake_case     :    torch.FloatTensor      ,					__snake_case     :    Union[torch.Tensor, float, int]      ,					__snake_case     :    Optional[torch.Tensor] = None      ,					__snake_case     :    bool = True      ,					):
         # 0. center input if necessary
         if self.config.center_input_sample:
             UpperCAmelCase_											=  2 * sample - 1.0

         # 1. time
         UpperCAmelCase_											=  timestep
         if not torch.is_tensor(__snake_case       ):
             UpperCAmelCase_											=  torch.tensor([timesteps]      ,					dtype=torch.long      ,					device=sample.device       )
         elif torch.is_tensor(__snake_case       ) and len(timesteps.shape       ) == 0:
             UpperCAmelCase_											=  timesteps[None].to(sample.device       )

         # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
         UpperCAmelCase_											=  timesteps * torch.ones(sample.shape[0]      ,					dtype=timesteps.dtype      ,					device=timesteps.device       )

         UpperCAmelCase_											=  self.time_proj(__snake_case       )

         # timesteps does not contain any weights and will always return f32 tensors
         # but time_embedding might actually be running in fp16. so we need to cast here.
         # there might be better ways to encapsulate this.
         UpperCAmelCase_											=  t_emb.to(dtype=self.dtype       )
         UpperCAmelCase_											=  self.time_embedding(__snake_case       )

         if self.class_embedding is not None:
             if class_labels is None:
                 raise ValueError('''class_labels should be provided when doing class conditioning'''       )

             if self.config.class_embed_type == "timestep":
                 UpperCAmelCase_											=  self.time_proj(__snake_case       )

             UpperCAmelCase_											=  self.class_embedding(__snake_case       ).to(dtype=self.dtype       )
             UpperCAmelCase_											=  emb + class_emb

         # 2. pre-process
         UpperCAmelCase_											=  sample
         UpperCAmelCase_											=  self.conv_in(__snake_case       )

         # 3. down
         UpperCAmelCase_											=  (sample,)
         for downsample_block in self.down_blocks:
             if hasattr(__snake_case      ,					'''skip_conv'''       ):
                 UpperCAmelCase_				,							UpperCAmelCase_				,							UpperCAmelCase_											=  downsample_block(
                     hidden_states=__snake_case      ,					temb=__snake_case      ,					skip_sample=__snake_case       )
             else:
                 UpperCAmelCase_				,							UpperCAmelCase_											=  downsample_block(hidden_states=__snake_case      ,					temb=__snake_case       )

             down_block_res_samples += res_samples

         # 4. mid
         UpperCAmelCase_											=  self.mid_block(__snake_case      ,					__snake_case       )

         # 5. up
         UpperCAmelCase_											=  None
         for upsample_block in self.up_blocks:
             UpperCAmelCase_											=  down_block_res_samples[-len(upsample_block.resnets       ) :]
             UpperCAmelCase_											=  down_block_res_samples[: -len(upsample_block.resnets       )]

             if hasattr(__snake_case      ,					'''skip_conv'''       ):
                 UpperCAmelCase_				,							UpperCAmelCase_											=  upsample_block(__snake_case      ,					__snake_case      ,					__snake_case      ,					__snake_case       )
             else:
                 UpperCAmelCase_											=  upsample_block(__snake_case      ,					__snake_case      ,					__snake_case       )

        # 6. post-process
         UpperCAmelCase_											=  self.conv_norm_out(__snake_case       )
         UpperCAmelCase_											=  self.conv_act(__snake_case       )
         UpperCAmelCase_											=  self.conv_out(__snake_case       )

         if skip_sample is not None:
             sample += skip_sample

         if self.config.time_embedding_type == "fourier":
             UpperCAmelCase_											=  timesteps.reshape((sample.shape[0], *([1] * len(sample.shape[1:]       )))       )
             UpperCAmelCase_											=  sample / timesteps

         if not return_dict:
             return (sample,)

         return UNetaDOutput(sample=__snake_case       )





