


'''simple docstring'''






import pickle
import shutil
import tempfile
import unittest

from transformers import SPIECE_UNDERLINE, XLMRobertaTokenizer, XLMRobertaTokenizerFast
from transformers.testing_utils import get_tests_dir, require_sentencepiece, require_tokenizers, slow
from transformers.utils import cached_property

from ...test_tokenization_common import TokenizerTesterMixin


__A :  str    							=				get_tests_dir("fixtures/test_sentencepiece.model")






@require_sentencepiece
@require_tokenizers
class        __snake_case	(       a__     ,unittest.TestCase):







    """simple docstring"""



    lowercase       			= XLMRobertaTokenizer
    lowercase       			= XLMRobertaTokenizerFast
    lowercase       			= True
    lowercase       			= True

    def __lowercase       (       self   :  Tuple		)       ->  List[Any]:
      super().setUp()

      # We have a SentencePiece fixture for testing
      lowerCAmelCase_      :				Tuple      =       XLMRobertaTokenizer(lowerCamelCase   ,						keep_accents=lowerCamelCase		)
      tokenizer.save_pretrained(self.tmpdirname		)

    def __lowercase       (       self   :  Optional[int]		)       ->  List[Any]:
      lowerCAmelCase_      :				Dict      =       "<pad>"
      lowerCAmelCase_      :				Any      =       1

      self.assertEqual(self.get_tokenizer()._convert_token_to_id(lowerCamelCase		)   ,						lowerCamelCase		)
      self.assertEqual(self.get_tokenizer()._convert_id_to_token(lowerCamelCase		)   ,						lowerCamelCase		)

    def __lowercase       (       self   :  Union[str, Any]		)       ->  List[str]:
      lowerCAmelCase_      :				Optional[int]      =       list(self.get_tokenizer().get_vocab().keys()		)

      self.assertEqual(vocab_keys[0]   ,						"""<s>"""		)
      self.assertEqual(vocab_keys[1]   ,						"""<pad>"""		)
      self.assertEqual(vocab_keys[-1]   ,						"""<mask>"""		)
      self.assertEqual(len(lowerCamelCase		)   ,						10_02		)

    def __lowercase       (       self   :  List[Any]		)       ->  Optional[int]:
      self.assertEqual(self.get_tokenizer().vocab_size   ,						10_02		)

    def __lowercase       (       self   :  Union[str, Any]		)       ->  Any:
      lowerCAmelCase_      :				Any      =       XLMRobertaTokenizer(lowerCamelCase   ,						keep_accents=lowerCamelCase		)

      lowerCAmelCase_      :				Dict      =       tokenizer.tokenize("""This is a test"""		)
      self.assertListEqual(lowerCamelCase   ,						["""▁This""", """▁is""", """▁a""", """▁t""", """est"""]		)

      self.assertListEqual(
          tokenizer.convert_tokens_to_ids(lowerCamelCase		)   ,						[value + tokenizer.fairseq_offset for value in [2_85, 46, 10, 1_70, 3_82]]   ,						)

      lowerCAmelCase_      :				Optional[Any]      =       tokenizer.tokenize("""I was born in 92000, and this is falsé."""		)
      self.assertListEqual(
          lowerCamelCase   ,						[
              SPIECE_UNDERLINE + """I""",
              SPIECE_UNDERLINE + """was""",
              SPIECE_UNDERLINE + """b""",
              """or""",
              """n""",
              SPIECE_UNDERLINE + """in""",
              SPIECE_UNDERLINE + """""",
              """9""",
              """2""",
              """0""",
              """0""",
              """0""",
              """,""",
              SPIECE_UNDERLINE + """and""",
              SPIECE_UNDERLINE + """this""",
              SPIECE_UNDERLINE + """is""",
              SPIECE_UNDERLINE + """f""",
              """al""",
              """s""",
              """é""",
              """.""",
          ]   ,						)
      lowerCAmelCase_      :				Optional[Any]      =       tokenizer.convert_tokens_to_ids(lowerCamelCase		)
      self.assertListEqual(
          lowerCamelCase   ,						[
              value + tokenizer.fairseq_offset
              for value in [8, 21, 84, 55, 24, 19, 7, 2, 6_02, 3_47, 3_47, 3_47, 3, 12, 66, 46, 72, 80, 6, 2, 4]
              #                                       ^ unk: 2 + 1 = 3                  unk: 2 + 1 = 3 ^
          ]   ,						)

      lowerCAmelCase_      :				Union[str, Any]      =       tokenizer.convert_ids_to_tokens(lowerCamelCase		)
      self.assertListEqual(
          lowerCamelCase   ,						[
              SPIECE_UNDERLINE + """I""",
              SPIECE_UNDERLINE + """was""",
              SPIECE_UNDERLINE + """b""",
              """or""",
              """n""",
              SPIECE_UNDERLINE + """in""",
              SPIECE_UNDERLINE + """""",
              """<unk>""",
              """2""",
              """0""",
              """0""",
              """0""",
              """,""",
              SPIECE_UNDERLINE + """and""",
              SPIECE_UNDERLINE + """this""",
              SPIECE_UNDERLINE + """is""",
              SPIECE_UNDERLINE + """f""",
              """al""",
              """s""",
              """<unk>""",
              """.""",
          ]   ,						)

    def __lowercase       (       self   :  Optional[Any]		)       ->  Dict:
      if not self.test_slow_tokenizer:
        # as we don't have a slow version, we can't compare the outputs between slow and fast versions
        return

      lowerCAmelCase_      :				Any      =       (self.rust_tokenizer_class, "hf-internal-testing/tiny-xlm-roberta", {})
      for tokenizer, pretrained_name, kwargs in self.tokenizers_list:
        with self.subTest(F'{tokenizer.__class__.__name__} ({pretrained_name})'		):
          lowerCAmelCase_      :				Optional[int]      =       self.rust_tokenizer_class.from_pretrained(lowerCamelCase   ,						**lowerCamelCase		)
          lowerCAmelCase_      :				int      =       self.tokenizer_class.from_pretrained(lowerCamelCase   ,						**lowerCamelCase		)

          lowerCAmelCase_      :				Optional[Any]      =       tempfile.mkdtemp()

          lowerCAmelCase_      :				Optional[Any]      =       tokenizer_r.save_pretrained(lowerCamelCase		)
          lowerCAmelCase_      :				Optional[Any]      =       tokenizer_p.save_pretrained(lowerCamelCase		)

          # Checks it save with the same files + the tokenizer.json file for the fast one
          self.assertTrue(any("""tokenizer.json""" in f for f in tokenizer_r_files		)		)
          lowerCAmelCase_      :				int      =       tuple(f for f in tokenizer_r_files if """tokenizer.json""" not in f		)
          self.assertSequenceEqual(lowerCamelCase   ,						lowerCamelCase		)

          # Checks everything loads correctly in the same way
          lowerCAmelCase_      :				str      =       tokenizer_r.from_pretrained(lowerCamelCase		)
          lowerCAmelCase_      :				Optional[Any]      =       tokenizer_p.from_pretrained(lowerCamelCase		)

          # Check special tokens are set accordingly on Rust and Python
          for key in tokenizer_pp.special_tokens_map:
            self.assertTrue(hasattr(lowerCamelCase   ,						lowerCamelCase		)		)
            # self.assertEqual(getattr(tokenizer_rp, key), getattr(tokenizer_pp, key))
            # self.assertEqual(getattr(tokenizer_rp, key + "_id"), getattr(tokenizer_pp, key + "_id"))

          shutil.rmtree(lowerCamelCase		)

          # Save tokenizer rust, legacy_format=True
          lowerCAmelCase_      :				List[str]      =       tempfile.mkdtemp()

          lowerCAmelCase_      :				Any      =       tokenizer_r.save_pretrained(lowerCamelCase   ,						legacy_format=lowerCamelCase		)
          lowerCAmelCase_      :				Union[str, Any]      =       tokenizer_p.save_pretrained(lowerCamelCase		)

          # Checks it save with the same files
          self.assertSequenceEqual(lowerCamelCase   ,						lowerCamelCase		)

          # Checks everything loads correctly in the same way
          lowerCAmelCase_      :				List[Any]      =       tokenizer_r.from_pretrained(lowerCamelCase		)
          lowerCAmelCase_      :				List[Any]      =       tokenizer_p.from_pretrained(lowerCamelCase		)

          # Check special tokens are set accordingly on Rust and Python
          for key in tokenizer_pp.special_tokens_map:
            self.assertTrue(hasattr(lowerCamelCase   ,						lowerCamelCase		)		)

          shutil.rmtree(lowerCamelCase		)

          # Save tokenizer rust, legacy_format=False
          lowerCAmelCase_      :				Optional[Any]      =       tempfile.mkdtemp()

          lowerCAmelCase_      :				int      =       tokenizer_r.save_pretrained(lowerCamelCase   ,						legacy_format=lowerCamelCase		)
          lowerCAmelCase_      :				Any      =       tokenizer_p.save_pretrained(lowerCamelCase		)

          # Checks it saved the tokenizer.json file
          self.assertTrue(any("""tokenizer.json""" in f for f in tokenizer_r_files		)		)

          # Checks everything loads correctly in the same way
          lowerCAmelCase_      :				Union[str, Any]      =       tokenizer_r.from_pretrained(lowerCamelCase		)
          lowerCAmelCase_      :				Dict      =       tokenizer_p.from_pretrained(lowerCamelCase		)

          # Check special tokens are set accordingly on Rust and Python
          for key in tokenizer_pp.special_tokens_map:
            self.assertTrue(hasattr(lowerCamelCase   ,						lowerCamelCase		)		)

          shutil.rmtree(lowerCamelCase		)

    @cached_property
    def __lowercase       (       self   :  Dict		)       ->  Dict:
      return XLMRobertaTokenizer.from_pretrained("""xlm-roberta-base"""		)

    def __lowercase       (       self   :  str		)       ->  List[Any]:
      with tempfile.NamedTemporaryFile() as f:
        shutil.copyfile(lowerCamelCase   ,						f.name		)
        lowerCAmelCase_      :				Optional[int]      =       XLMRobertaTokenizer(f.name   ,						keep_accents=lowerCamelCase		)
        lowerCAmelCase_      :				str      =       pickle.dumps(lowerCamelCase		)
      pickle.loads(lowerCamelCase		)

    def __lowercase       (       self   :  Dict		)       ->  Optional[Any]:
      if not self.test_rust_tokenizer:
        return

      lowerCAmelCase_      :				Union[str, Any]      =       self.get_tokenizer()
      lowerCAmelCase_      :				Optional[int]      =       self.get_rust_tokenizer()

      lowerCAmelCase_      :				int      =       "I was born in 92000, and this is falsé."

      lowerCAmelCase_      :				str      =       tokenizer.tokenize(lowerCamelCase		)
      lowerCAmelCase_      :				List[Any]      =       rust_tokenizer.tokenize(lowerCamelCase		)
      self.assertListEqual(lowerCamelCase   ,						lowerCamelCase		)

      lowerCAmelCase_      :				Dict      =       tokenizer.encode(lowerCamelCase   ,						add_special_tokens=lowerCamelCase		)
      lowerCAmelCase_      :				Optional[Any]      =       rust_tokenizer.encode(lowerCamelCase   ,						add_special_tokens=lowerCamelCase		)
      self.assertListEqual(lowerCamelCase   ,						lowerCamelCase		)

      lowerCAmelCase_      :				List[str]      =       self.get_rust_tokenizer()
      lowerCAmelCase_      :				Dict      =       tokenizer.encode(lowerCamelCase		)
      lowerCAmelCase_      :				Optional[int]      =       rust_tokenizer.encode(lowerCamelCase		)
      self.assertListEqual(lowerCamelCase   ,						lowerCamelCase		)

    @slow
    def __lowercase       (       self   :  Union[str, Any]		)       ->  int:
      lowerCAmelCase_      :				List[Any]      =       "Hello World!"
      lowerCAmelCase_      :				int      =       [0, 3_53_78, 66_61, 38, 2]
      # xlmr = torch.hub.load('pytorch/fairseq', 'xlmr.base')  # xlmr.large has same tokenizer
      # xlmr.eval()
      # xlmr.encode(symbols)

      self.assertListEqual(lowerCamelCase   ,						self.big_tokenizer.encode(lowerCamelCase		)		)

    @slow
    def __lowercase       (       self   :  List[str]		)       ->  Dict:
      lowerCAmelCase_      :				str      =       (
          "This is a very long text with a lot of weird characters, such as: . , ~ ? ( ) \" [ ] ! : - . Also we will"
          " add words that should not exsist and be tokenized to <unk>, such as saoneuhaoesuth"
      )
      lowerCAmelCase_      :				int      =       [
          0,
          32_93,
          83,
          10,
          45_52,
          49_89,
          79_86,
          6_78,
          10,
          59_15,
          1_11,
          17_94_59,
          12_48_50,
          4,
          60_44,
          2_37,
          12,
          6,
          5,
          6,
          4,
          67_80,
          7_05,
          15,
          13_88,
          44,
          3_78,
          1_01_14,
          7_11,
          1_52,
          20,
          6,
          5,
          2_23_76,
          6_42,
          12_21,
          1_51_90,
          3_41_53,
          4_50,
          56_08,
          9_59,
          11_19,
          5_77_02,
          1_36,
          1_86,
          47,
          10_98,
          2_93_67,
          47,
          # 4426, # What fairseq tokenizes from "<unk>": "_<"
          # 3678, # What fairseq tokenizes from "<unk>": "unk"
          # 2740, # What fairseq tokenizes from "<unk>": ">"
          3,  # What we tokenize from "<unk>": "<unk>"
          6,  # Residue from the tokenization: an extra sentencepiece underline
          4,
          60_44,
          2_37,
          62_84,
          5_09_01,
          5_28,
          31,
          90,
          34,
          9_27,
          2,
      ]
      # xlmr = torch.hub.load('pytorch/fairseq', 'xlmr.base')  # xlmr.large has same tokenizer
      # xlmr.eval()
      # xlmr.encode(symbols)

      self.assertListEqual(lowerCamelCase   ,						self.big_tokenizer.encode(lowerCamelCase		)		)





    @slow
    def __lowercase       (       self   :  Dict		)       ->  Tuple:
      # fmt: off
      lowerCAmelCase_      :				Optional[int]      =       {"input_ids": [[0, 1_10_62, 8_27_72, 7, 15, 8_27_72, 5_38, 5_15_29, 2_37, 1_71_98, 12_90, 2_06, 9, 21_51_75, 13_14, 1_36, 1_71_98, 12_90, 2_06, 9, 5_63_59, 42, 12_20_09, 9, 1_64_66, 16, 8_73_44, 45_37, 9, 47_17, 7_83_81, 6, 15_99_58, 7, 15, 2_44_80, 6_18, 4, 5_27, 2_26_93, 54_28, 4, 27_77, 2_44_80, 98_74, 4, 4_35_23, 5_94, 4, 8_03, 1_83_92, 3_31_89, 18, 4, 4_35_23, 2_44_47, 1_23_99, 1_00, 2_49_55, 8_36_58, 96_26, 14_40_57, 15, 8_39, 2_23_35, 16, 1_36, 2_49_55, 8_36_58, 8_34_79, 15, 3_91_02, 7_24, 16, 6_78, 6_45, 27_89, 13_28, 45_89, 42, 12_20_09, 11_57_74, 23, 8_05, 13_28, 4_68_76, 7, 1_36, 5_38_94, 19_40, 4_22_27, 4_11_59, 1_77_21, 8_23, 4_25, 4, 2_75_12, 9_87_22, 2_06, 1_36, 55_31, 49_70, 9_19, 1_73_36, 5, 2], [0, 2_00_80, 6_18, 83, 8_27_75, 47, 4_79, 9, 15_17, 73, 5_38_94, 3_33, 8_05_81, 11_01_17, 1_88_11, 52_56, 12_95, 51, 15_25_26, 2_97, 79_86, 3_90, 12_44_16, 5_38, 3_54_31, 2_14, 98, 1_50_44, 2_57_37, 1_36, 71_08, 4_37_01, 23, 7_56, 13_53_55, 7, 5, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [0, 5_81, 6_37_73, 11_94_55, 6, 14_77_97, 8_82_03, 7, 6_45, 70, 21, 32_85, 1_02_69, 5, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]], "attention_mask": [[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]}  # noqa: E501
      # fmt: on

      self.tokenizer_integration_test_util(
          expected_encoding=lowerCamelCase   ,						model_name="""xlm-roberta-base"""   ,						revision="""d9d8a8ea5eb94b1c6654ae9249df7793cd2933d3"""   ,						)







'''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        __snake_case	(       _SCREAMING_SNAKE_CASE):







    """simple docstring"""



    lowercase       			= 42







class        __snake_case	(       _SCREAMING_SNAKE_CASE     ,_SCREAMING_SNAKE_CASE):







    """simple docstring"""



    @register_to_config
    def __init__(       self   :  Optional[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   ,						)       ->  List[Any]:
      super().__init__()
      lowerCAmelCase_      :				Optional[Any]      =       sample_size

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

      if use_timestep_embedding:
        lowerCAmelCase_      :				List[str]      =       block_out_channels[0] * 4
        lowerCAmelCase_      :				List[str]      =       TimestepEmbedding(
            in_channels=lowerCamelCase   ,						time_embed_dim=lowerCamelCase   ,						act_fn=lowerCamelCase   ,						out_dim=block_out_channels[0]   ,						)

      lowerCAmelCase_      :				str      =       nn.ModuleList([]		)
      lowerCAmelCase_      :				Tuple      =       None
      lowerCAmelCase_      :				str      =       nn.ModuleList([]		)
      lowerCAmelCase_      :				List[str]      =       None

      # down
      lowerCAmelCase_      :				List[str]      =       in_channels
      for i, down_block_type in enumerate(lowerCamelCase		):
        lowerCAmelCase_      :				Optional[Any]      =       output_channel
        lowerCAmelCase_      :				str      =       block_out_channels[i]

        if i == 0:
          input_channel += extra_in_channels

        lowerCAmelCase_      :				int      =       i == len(lowerCamelCase		) - 1

        lowerCAmelCase_      :				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
      lowerCAmelCase_      :				Dict      =       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
      lowerCAmelCase_      :				List[str]      =       list(reversed(lowerCamelCase		)		)
      lowerCAmelCase_      :				str      =       reversed_block_out_channels[0]
      if out_block_type is None:
        lowerCAmelCase_      :				Optional[int]      =       out_channels
      else:
        lowerCAmelCase_      :				Optional[Any]      =       block_out_channels[0]

      for i, up_block_type in enumerate(lowerCamelCase		):
        lowerCAmelCase_      :				Any      =       output_channel
        lowerCAmelCase_      :				str      =       (
            reversed_block_out_channels[i + 1] if i < len(lowerCamelCase		) - 1 else final_upsample_channels
        )

        lowerCAmelCase_      :				int      =       i == len(lowerCamelCase		) - 1

        lowerCAmelCase_      :				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		)
        lowerCAmelCase_      :				Optional[int]      =       output_channel

      # out
      lowerCAmelCase_      :				str      =       norm_num_groups if norm_num_groups is not None else min(block_out_channels[0] // 4   ,						32		)
      lowerCAmelCase_      :				Union[str, Any]      =       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 __lowercase       (       self   :  List[Any]   ,						lowerCamelCase   :  torch.FloatTensor   ,						lowerCamelCase   :  Union[torch.Tensor, float, int]   ,						lowerCamelCase   :  bool = True   ,						)       ->  Union[UNetaDOutput, Tuple]:
      lowerCAmelCase_      :				Optional[Any]      =       timestep
      if not torch.is_tensor(lowerCamelCase		):
        lowerCAmelCase_      :				Any      =       torch.tensor([timesteps]   ,						dtype=torch.long   ,						device=sample.device		)
      elif torch.is_tensor(lowerCamelCase		) and len(timesteps.shape		) == 0:
        lowerCAmelCase_      :				Dict      =       timesteps[None].to(sample.device		)

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

      # 2. down
      lowerCAmelCase_      :				Optional[Any]      =       ()
      for downsample_block in self.down_blocks:
        lowerCAmelCase_,							lowerCAmelCase_      :				Any      =       downsample_block(hidden_states=lowerCamelCase   ,						temb=lowerCamelCase		)
        down_block_res_samples += res_samples

      # 3. mid
      if self.mid_block:
        lowerCAmelCase_      :				int      =       self.mid_block(lowerCamelCase   ,						lowerCamelCase		)

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

      # 5. post-process
      if self.out_block:
        lowerCAmelCase_      :				Union[str, Any]      =       self.out_block(lowerCamelCase   ,						lowerCamelCase		)

      if not return_dict:
        return (sample,)

      return UNetaDOutput(sample=lowerCamelCase		)



