
"""simple docstring"""




import unittest

from transformers import MarkupLMConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device
from transformers.utils import cached_property

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
    import torch

    from transformers import (
        MarkupLMForQuestionAnswering,
        MarkupLMForSequenceClassification,
        MarkupLMForTokenClassification,
        MarkupLMModel,
    )

# TODO check dependencies
from transformers import MarkupLMFeatureExtractor, MarkupLMProcessor, MarkupLMTokenizer







class 	_lowercase    :


     '''simple docstring'''





     def __init__(   self			,      __UpperCamelCase			,      __UpperCamelCase=13			,      __UpperCamelCase=7			,      __UpperCamelCase=True			,      __UpperCamelCase=True			,      __UpperCamelCase=True			,      __UpperCamelCase=True			,      __UpperCamelCase=99			,      __UpperCamelCase=32			,      __UpperCamelCase=5			,      __UpperCamelCase=4			,      __UpperCamelCase=37			,      __UpperCamelCase="gelu"			,      __UpperCamelCase=0.1			,      __UpperCamelCase=0.1			,      __UpperCamelCase=5_12			,      __UpperCamelCase=16			,      __UpperCamelCase=2			,      __UpperCamelCase=0.02			,      __UpperCamelCase=3			,      __UpperCamelCase=None			,      __UpperCamelCase=20			,      __UpperCamelCase=30			,      __UpperCamelCase=2			,      __UpperCamelCase=2			,      __UpperCamelCase=10			,      )->       Dict:
          UpperCAmelCase__				:			Optional[Any]							=    parent
          UpperCAmelCase__				:			str							=    batch_size
          UpperCAmelCase__				:			Tuple							=    seq_length
          UpperCAmelCase__				:			List[Any]							=    is_training
          UpperCAmelCase__				:			Any							=    use_input_mask
          UpperCAmelCase__				:			Optional[Any]							=    use_token_type_ids
          UpperCAmelCase__				:			Optional[Any]							=    use_labels
          UpperCAmelCase__				:			Dict							=    vocab_size
          UpperCAmelCase__				:			List[str]							=    hidden_size
          UpperCAmelCase__				:			List[Any]							=    num_hidden_layers
          UpperCAmelCase__				:			Optional[int]							=    num_attention_heads
          UpperCAmelCase__				:			List[Any]							=    intermediate_size
          UpperCAmelCase__				:			Optional[Any]							=    hidden_act
          UpperCAmelCase__				:			str							=    hidden_dropout_prob
          UpperCAmelCase__				:			Any							=    attention_probs_dropout_prob
          UpperCAmelCase__				:			int							=    max_position_embeddings
          UpperCAmelCase__				:			List[Any]							=    type_vocab_size
          UpperCAmelCase__				:			List[Any]							=    type_sequence_label_size
          UpperCAmelCase__				:			List[Any]							=    initializer_range
          UpperCAmelCase__				:			Dict							=    num_labels
          UpperCAmelCase__				:			Optional[int]							=    scope
          UpperCAmelCase__				:			List[Any]							=    max_xpath_tag_unit_embeddings
          UpperCAmelCase__				:			str							=    max_xpath_subs_unit_embeddings
          UpperCAmelCase__				:			Any							=    tag_pad_id
          UpperCAmelCase__				:			Optional[int]							=    subs_pad_id
          UpperCAmelCase__				:			Tuple							=    max_depth



     def 			lowerCAmelCase__							(   self	)->       str:
          UpperCAmelCase__				:			Any							=    ids_tensor([self.batch_size, self.seq_length]			,      self.vocab_size	)

          UpperCAmelCase__				:			Dict							=    ids_tensor(
              [self.batch_size, self.seq_length, self.max_depth]			,      self.max_xpath_tag_unit_embeddings	)

          UpperCAmelCase__				:			Dict							=    ids_tensor(
              [self.batch_size, self.seq_length, self.max_depth]			,      self.max_xpath_subs_unit_embeddings	)

          UpperCAmelCase__				:			Any							=    None
          if self.use_input_mask:
               UpperCAmelCase__				:			List[str]							=    ids_tensor([self.batch_size, self.seq_length]			,      vocab_size=2	)

          UpperCAmelCase__				:			Any							=    None
          if self.use_token_type_ids:
               UpperCAmelCase__				:			Union[str, Any]							=    ids_tensor([self.batch_size, self.seq_length]			,      self.type_vocab_size	)

          UpperCAmelCase__				:			Any							=    None
          UpperCAmelCase__				:			Any							=    None
          if self.use_labels:
               UpperCAmelCase__				:			Union[str, Any]							=    ids_tensor([self.batch_size]			,      self.type_sequence_label_size	)
               UpperCAmelCase__				:			int							=    ids_tensor([self.batch_size, self.seq_length]			,      self.num_labels	)

          UpperCAmelCase__				:			str							=    self.get_config()

          return (
              config,
              input_ids,
              xpath_tags_seq,
              xpath_subs_seq,
              token_type_ids,
              input_mask,
              sequence_labels,
              token_labels,
          )



     def 			lowerCAmelCase__							(   self	)->       str:
          return MarkupLMConfig(
              vocab_size=self.vocab_size			,      hidden_size=self.hidden_size			,      num_hidden_layers=self.num_hidden_layers			,      num_attention_heads=self.num_attention_heads			,      intermediate_size=self.intermediate_size			,      hidden_act=self.hidden_act			,      hidden_dropout_prob=self.hidden_dropout_prob			,      attention_probs_dropout_prob=self.attention_probs_dropout_prob			,      max_position_embeddings=self.max_position_embeddings			,      type_vocab_size=self.type_vocab_size			,      initializer_range=self.initializer_range			,      max_xpath_tag_unit_embeddings=self.max_xpath_tag_unit_embeddings			,      max_xpath_subs_unit_embeddings=self.max_xpath_subs_unit_embeddings			,      tag_pad_id=self.tag_pad_id			,      subs_pad_id=self.subs_pad_id			,      max_depth=self.max_depth			,      )



     def 			lowerCAmelCase__							(   self			,      __UpperCamelCase			,      __UpperCamelCase			,      __UpperCamelCase			,      __UpperCamelCase			,      __UpperCamelCase			,      __UpperCamelCase			,      __UpperCamelCase			,      __UpperCamelCase			,      )->       List[Any]:
          UpperCAmelCase__				:			int							=    MarkupLMModel(config=__UpperCamelCase	)
          model.to(__UpperCamelCase	)
          model.eval()
          print("Configs:"			,      model.config.tag_pad_id			,      model.config.subs_pad_id	)
          UpperCAmelCase__				:			Tuple							=    model(__UpperCamelCase			,      attention_mask=__UpperCamelCase			,      token_type_ids=__UpperCamelCase	)
          UpperCAmelCase__				:			int							=    model(__UpperCamelCase			,      token_type_ids=__UpperCamelCase	)
          UpperCAmelCase__				:			Any							=    model(__UpperCamelCase	)
          self.parent.assertEqual(result.last_hidden_state.shape			,      (self.batch_size, self.seq_length, self.hidden_size)	)
          self.parent.assertEqual(result.pooler_output.shape			,      (self.batch_size, self.hidden_size)	)



     def 			lowerCAmelCase__							(   self			,      __UpperCamelCase			,      __UpperCamelCase			,      __UpperCamelCase			,      __UpperCamelCase			,      __UpperCamelCase			,      __UpperCamelCase			,      __UpperCamelCase			,      __UpperCamelCase			,      )->       List[Any]:
          UpperCAmelCase__				:			Optional[int]							=    self.num_labels
          UpperCAmelCase__				:			Tuple							=    MarkupLMForSequenceClassification(__UpperCamelCase	)
          model.to(__UpperCamelCase	)
          model.eval()
          UpperCAmelCase__				:			str							=    model(
              __UpperCamelCase			,      xpath_tags_seq=__UpperCamelCase			,      xpath_subs_seq=__UpperCamelCase			,      attention_mask=__UpperCamelCase			,      token_type_ids=__UpperCamelCase			,      labels=__UpperCamelCase			,      )
          self.parent.assertEqual(result.logits.shape			,      (self.batch_size, self.num_labels)	)



     def 			lowerCAmelCase__							(   self			,      __UpperCamelCase			,      __UpperCamelCase			,      __UpperCamelCase			,      __UpperCamelCase			,      __UpperCamelCase			,      __UpperCamelCase			,      __UpperCamelCase			,      __UpperCamelCase			,      )->       Any:
          UpperCAmelCase__				:			Union[str, Any]							=    self.num_labels
          UpperCAmelCase__				:			str							=    MarkupLMForTokenClassification(config=__UpperCamelCase	)
          model.to(__UpperCamelCase	)
          model.eval()
          UpperCAmelCase__				:			Dict							=    model(
              __UpperCamelCase			,      xpath_tags_seq=__UpperCamelCase			,      xpath_subs_seq=__UpperCamelCase			,      attention_mask=__UpperCamelCase			,      token_type_ids=__UpperCamelCase			,      labels=__UpperCamelCase			,      )
          self.parent.assertEqual(result.logits.shape			,      (self.batch_size, self.seq_length, self.num_labels)	)



     def 			lowerCAmelCase__							(   self			,      __UpperCamelCase			,      __UpperCamelCase			,      __UpperCamelCase			,      __UpperCamelCase			,      __UpperCamelCase			,      __UpperCamelCase			,      __UpperCamelCase			,      __UpperCamelCase			,      )->       Dict:
          UpperCAmelCase__				:			Dict							=    MarkupLMForQuestionAnswering(config=__UpperCamelCase	)
          model.to(__UpperCamelCase	)
          model.eval()
          UpperCAmelCase__				:			Tuple							=    model(
              __UpperCamelCase			,      xpath_tags_seq=__UpperCamelCase			,      xpath_subs_seq=__UpperCamelCase			,      attention_mask=__UpperCamelCase			,      token_type_ids=__UpperCamelCase			,      start_positions=__UpperCamelCase			,      end_positions=__UpperCamelCase			,      )
          self.parent.assertEqual(result.start_logits.shape			,      (self.batch_size, self.seq_length)	)
          self.parent.assertEqual(result.end_logits.shape			,      (self.batch_size, self.seq_length)	)



     def 			lowerCAmelCase__							(   self	)->       Optional[int]:
          UpperCAmelCase__				:			Optional[int]							=    self.prepare_config_and_inputs()
          (
              (
              UpperCAmelCase__
          )				,						(
              UpperCAmelCase__
          )				,						(
              UpperCAmelCase__
          )				,						(
              UpperCAmelCase__
          )				,						(
              UpperCAmelCase__
          )				,						(
              UpperCAmelCase__
          )				,						(
              UpperCAmelCase__
          )				,						(
              UpperCAmelCase__
          )				,						
          )				:			List[str]							=    config_and_inputs
          UpperCAmelCase__				:			Any							=    {
              "input_ids": input_ids,
              "xpath_tags_seq": xpath_tags_seq,
              "xpath_subs_seq": xpath_subs_seq,
              "token_type_ids": token_type_ids,
              "attention_mask": input_mask,
          }
          return config, inputs_dict







@require_torch
class 	_lowercase    (							lowerCAmelCase_						, lowerCAmelCase_						, unittest.TestCase					):


     '''simple docstring'''





     _A		    =				(
         (
             MarkupLMModel,
             MarkupLMForSequenceClassification,
             MarkupLMForTokenClassification,
             MarkupLMForQuestionAnswering,
         )
         if is_torch_available()
         else None
     )
     _A		    =				(
         {
             'feature-extraction': MarkupLMModel,
             'question-answering': MarkupLMForQuestionAnswering,
             'text-classification': MarkupLMForSequenceClassification,
             'token-classification': MarkupLMForTokenClassification,
             'zero-shot': MarkupLMForSequenceClassification,
         }
         if is_torch_available()
         else {}
     )



     def 			lowerCAmelCase__							(   self			,      __UpperCamelCase			,      __UpperCamelCase			,      __UpperCamelCase			,      __UpperCamelCase			,      __UpperCamelCase	)->       Union[str, Any]:
          # ValueError: Nodes must be of type `List[str]` (single pretokenized example), or `List[List[str]]`
          # (batch of pretokenized examples).
          return True



     def 			lowerCAmelCase__							(   self	)->       Tuple:
          UpperCAmelCase__				:			str							=    MarkupLMModelTester(self	)
          UpperCAmelCase__				:			str							=    ConfigTester(self			,      config_class=__UpperCamelCase			,      hidden_size=37	)



     def 			lowerCAmelCase__							(   self	)->       int:
          self.config_tester.run_common_tests()



     def 			lowerCAmelCase__							(   self	)->       int:
          UpperCAmelCase__				:			List[Any]							=    self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_model(*__UpperCamelCase	)



     def 			lowerCAmelCase__							(   self	)->       Optional[Any]:
          UpperCAmelCase__				:			Union[str, Any]							=    self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_sequence_classification(*__UpperCamelCase	)



     def 			lowerCAmelCase__							(   self	)->       str:
          UpperCAmelCase__				:			List[Any]							=    self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_token_classification(*__UpperCamelCase	)



     def 			lowerCAmelCase__							(   self	)->       Tuple:
          UpperCAmelCase__				:			Dict							=    self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_question_answering(*__UpperCamelCase	)





def 				a__			(					):






     '''simple docstring'''
     UpperCAmelCase__				:			Optional[Any]							=    "\n    <!DOCTYPE html>\n    <html>\n    <head>\n    <title>Page Title</title>\n    </head>\n    <body>\n\n    <h1>This is a Heading</h1>\n    <p>This is a paragraph.</p>\n\n    </body>\n    </html>\n    "

     return html_string







@require_torch
class 	_lowercase    (							unittest.TestCase					):


     '''simple docstring'''





     @cached_property
     def 			lowerCAmelCase__							(   self	)->       int:
          # TODO use from_pretrained here
          UpperCAmelCase__				:			Dict							=    MarkupLMFeatureExtractor()
          UpperCAmelCase__				:			List[str]							=    MarkupLMTokenizer.from_pretrained("microsoft/markuplm-base"	)

          return MarkupLMProcessor(__UpperCamelCase			,      __UpperCamelCase	)



     @slow
     def 			lowerCAmelCase__							(   self	)->       Tuple:
          UpperCAmelCase__				:			str							=    MarkupLMModel.from_pretrained("microsoft/markuplm-base"	).to(__UpperCamelCase	)

          UpperCAmelCase__				:			Union[str, Any]							=    self.default_processor

          UpperCAmelCase__				:			Optional[Any]							=    processor(prepare_html_string()			,      return_tensors="pt"	)
          UpperCAmelCase__				:			Optional[Any]							=    inputs.to(__UpperCamelCase	)

          # forward pass
          with torch.no_grad():
               UpperCAmelCase__				:			Union[str, Any]							=    model(**__UpperCamelCase	)

          # verify the last hidden states
          UpperCAmelCase__				:			List[str]							=    torch.Size([1, 14, 7_68]	)
          self.assertEqual(outputs.last_hidden_state.shape			,      __UpperCamelCase	)

          UpperCAmelCase__				:			str							=    torch.tensor(
              [[0.0675, -0.0052, 0.5001], [-0.2281, 0.0802, 0.2192], [-0.0583, -0.3311, 0.1185]]	).to(__UpperCamelCase	)

          self.assertTrue(torch.allclose(outputs.last_hidden_state[0, :3, :3]			,      __UpperCamelCase			,      atol=1E-4	)	)





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







@dataclass
class 	_lowercase    (							lowerCAmelCase_					):


     '''simple docstring'''





     _A		    =				42







class 	_lowercase    (							lowerCAmelCase_						, lowerCAmelCase_					):


     '''simple docstring'''





     @register_to_config
     def __init__(   self			,      __UpperCamelCase = None			,      __UpperCamelCase = 3			,      __UpperCamelCase = 3			,      __UpperCamelCase = False			,      __UpperCamelCase = "positional"			,      __UpperCamelCase = 0			,      __UpperCamelCase = True			,      __UpperCamelCase = ("DownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D")			,      __UpperCamelCase = ("AttnUpBlock2D", "AttnUpBlock2D", "AttnUpBlock2D", "UpBlock2D")			,      __UpperCamelCase = (2_24, 4_48, 6_72, 8_96)			,      __UpperCamelCase = 2			,      __UpperCamelCase = 1			,      __UpperCamelCase = 1			,      __UpperCamelCase = "conv"			,      __UpperCamelCase = "conv"			,      __UpperCamelCase = "silu"			,      __UpperCamelCase = 8			,      __UpperCamelCase = 32			,      __UpperCamelCase = 1E-5			,      __UpperCamelCase = "default"			,      __UpperCamelCase = True			,      __UpperCamelCase = None			,      __UpperCamelCase = None			,      )->       Any:
          super().__init__()

          UpperCAmelCase__				:			Optional[int]							=    sample_size
          UpperCAmelCase__				:			Any							=    block_out_channels[0] * 4

          # Check inputs
          if len(__UpperCamelCase	) != len(__UpperCamelCase	):
               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(__UpperCamelCase	) != len(__UpperCamelCase	):
               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__				:			List[Any]							=    nn.Convad(__UpperCamelCase			,      block_out_channels[0]			,      kernel_size=3			,      padding=(1, 1)	)

          # time
          if time_embedding_type == "fourier":
               UpperCAmelCase__				:			int							=    GaussianFourierProjection(embedding_size=block_out_channels[0]			,      scale=16	)
               UpperCAmelCase__				:			Dict							=    2 * block_out_channels[0]
          elif time_embedding_type == "positional":
               UpperCAmelCase__				:			List[str]							=    Timesteps(block_out_channels[0]			,      __UpperCamelCase			,      __UpperCamelCase	)
               UpperCAmelCase__				:			Any							=    block_out_channels[0]

          UpperCAmelCase__				:			List[str]							=    TimestepEmbedding(__UpperCamelCase			,      __UpperCamelCase	)

          # class embedding
          if class_embed_type is None and num_class_embeds is not None:
               UpperCAmelCase__				:			Optional[Any]							=    nn.Embedding(__UpperCamelCase			,      __UpperCamelCase	)
          elif class_embed_type == "timestep":
               UpperCAmelCase__				:			str							=    TimestepEmbedding(__UpperCamelCase			,      __UpperCamelCase	)
          elif class_embed_type == "identity":
               UpperCAmelCase__				:			Union[str, Any]							=    nn.Identity(__UpperCamelCase			,      __UpperCamelCase	)
          else:
               UpperCAmelCase__				:			Optional[int]							=    None

          UpperCAmelCase__				:			str							=    nn.ModuleList([]	)
          UpperCAmelCase__				:			Any							=    None
          UpperCAmelCase__				:			Tuple							=    nn.ModuleList([]	)

          # down
          UpperCAmelCase__				:			Any							=    block_out_channels[0]
          for i, down_block_type in enumerate(__UpperCamelCase	):
               UpperCAmelCase__				:			List[Any]							=    output_channel
               UpperCAmelCase__				:			Optional[int]							=    block_out_channels[i]
               UpperCAmelCase__				:			Any							=    i == len(__UpperCamelCase	) - 1

               UpperCAmelCase__				:			List[Any]							=    get_down_block(
                   __UpperCamelCase			,      num_layers=__UpperCamelCase			,      in_channels=__UpperCamelCase			,      out_channels=__UpperCamelCase			,      temb_channels=__UpperCamelCase			,      add_downsample=not is_final_block			,      resnet_eps=__UpperCamelCase			,      resnet_act_fn=__UpperCamelCase			,      resnet_groups=__UpperCamelCase			,      attention_head_dim=attention_head_dim if attention_head_dim is not None else output_channel			,      downsample_padding=__UpperCamelCase			,      resnet_time_scale_shift=__UpperCamelCase			,      downsample_type=__UpperCamelCase			,      )
               self.down_blocks.append(__UpperCamelCase	)

          # mid
          UpperCAmelCase__				:			List[str]							=    UNetMidBlockaD(
              in_channels=block_out_channels[-1]			,      temb_channels=__UpperCamelCase			,      resnet_eps=__UpperCamelCase			,      resnet_act_fn=__UpperCamelCase			,      output_scale_factor=__UpperCamelCase			,      resnet_time_scale_shift=__UpperCamelCase			,      attention_head_dim=attention_head_dim if attention_head_dim is not None else block_out_channels[-1]			,      resnet_groups=__UpperCamelCase			,      add_attention=__UpperCamelCase			,      )

          # up
          UpperCAmelCase__				:			Tuple							=    list(reversed(__UpperCamelCase	)	)
          UpperCAmelCase__				:			Optional[Any]							=    reversed_block_out_channels[0]
          for i, up_block_type in enumerate(__UpperCamelCase	):
               UpperCAmelCase__				:			int							=    output_channel
               UpperCAmelCase__				:			List[Any]							=    reversed_block_out_channels[i]
               UpperCAmelCase__				:			Dict							=    reversed_block_out_channels[min(i + 1			,      len(__UpperCamelCase	) - 1	)]

               UpperCAmelCase__				:			int							=    i == len(__UpperCamelCase	) - 1

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

          # out
          UpperCAmelCase__				:			str							=    norm_num_groups if norm_num_groups is not None else min(block_out_channels[0] // 4			,      32	)
          UpperCAmelCase__				:			Tuple							=    nn.GroupNorm(num_channels=block_out_channels[0]			,      num_groups=__UpperCamelCase			,      eps=__UpperCamelCase	)
          UpperCAmelCase__				:			Any							=    nn.SiLU()
          UpperCAmelCase__				:			Optional[int]							=    nn.Convad(block_out_channels[0]			,      __UpperCamelCase			,      kernel_size=3			,      padding=1	)



     def 			lowerCAmelCase__							(   self			,      __UpperCamelCase			,      __UpperCamelCase			,      __UpperCamelCase = None			,      __UpperCamelCase = True			,      )->       Union[UNetaDOutput, Tuple]:
          # 0. center input if necessary
          if self.config.center_input_sample:
               UpperCAmelCase__				:			Any							=    2 * sample - 1.0

          # 1. time
          UpperCAmelCase__				:			Tuple							=    timestep
          if not torch.is_tensor(__UpperCamelCase	):
               UpperCAmelCase__				:			Optional[int]							=    torch.tensor([timesteps]			,      dtype=torch.long			,      device=sample.device	)
          elif torch.is_tensor(__UpperCamelCase	) and len(timesteps.shape	) == 0:
               UpperCAmelCase__				:			List[Any]							=    timesteps[None].to(sample.device	)

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

          UpperCAmelCase__				:			List[str]							=    self.time_proj(__UpperCamelCase	)

          # 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__				:			Dict							=    t_emb.to(dtype=self.dtype	)
          UpperCAmelCase__				:			Optional[int]							=    self.time_embedding(__UpperCamelCase	)

          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__				:			List[Any]							=    self.time_proj(__UpperCamelCase	)

               UpperCAmelCase__				:			List[str]							=    self.class_embedding(__UpperCamelCase	).to(dtype=self.dtype	)
               UpperCAmelCase__				:			Optional[int]							=    emb + class_emb

          # 2. pre-process
          UpperCAmelCase__				:			List[str]							=    sample
          UpperCAmelCase__				:			str							=    self.conv_in(__UpperCamelCase	)

          # 3. down
          UpperCAmelCase__				:			List[str]							=    (sample,)
          for downsample_block in self.down_blocks:
               if hasattr(__UpperCamelCase			,      "skip_conv"	):
                    UpperCAmelCase__				,						UpperCAmelCase__				,						UpperCAmelCase__				:			Optional[Any]							=    downsample_block(
                        hidden_states=__UpperCamelCase			,      temb=__UpperCamelCase			,      skip_sample=__UpperCamelCase	)
               else:
                    UpperCAmelCase__				,						UpperCAmelCase__				:			Optional[Any]							=    downsample_block(hidden_states=__UpperCamelCase			,      temb=__UpperCamelCase	)

               down_block_res_samples += res_samples

          # 4. mid
          UpperCAmelCase__				:			Union[str, Any]							=    self.mid_block(__UpperCamelCase			,      __UpperCamelCase	)

          # 5. up
          UpperCAmelCase__				:			Dict							=    None
          for upsample_block in self.up_blocks:
               UpperCAmelCase__				:			Union[str, Any]							=    down_block_res_samples[-len(upsample_block.resnets	) :]
               UpperCAmelCase__				:			Optional[Any]							=    down_block_res_samples[: -len(upsample_block.resnets	)]

               if hasattr(__UpperCamelCase			,      "skip_conv"	):
                    UpperCAmelCase__				,						UpperCAmelCase__				:			str							=    upsample_block(__UpperCamelCase			,      __UpperCamelCase			,      __UpperCamelCase			,      __UpperCamelCase	)
               else:
                    UpperCAmelCase__				:			Tuple							=    upsample_block(__UpperCamelCase			,      __UpperCamelCase			,      __UpperCamelCase	)

        # 6. post-process
          UpperCAmelCase__				:			List[Any]							=    self.conv_norm_out(__UpperCamelCase	)
          UpperCAmelCase__				:			Any							=    self.conv_act(__UpperCamelCase	)
          UpperCAmelCase__				:			Optional[Any]							=    self.conv_out(__UpperCamelCase	)

          if skip_sample is not None:
               sample += skip_sample

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

          if not return_dict:
               return (sample,)

          return UNetaDOutput(sample=__UpperCamelCase	)



