






import unittest

from transformers import (
    MODEL_FOR_CAUSAL_LM_MAPPING,
    TF_MODEL_FOR_CAUSAL_LM_MAPPING,
    TextGenerationPipeline,
    logging,
    pipeline,
)
from transformers.testing_utils import (
    CaptureLogger,
    is_pipeline_test,
    require_accelerate,
    require_tf,
    require_torch,
    require_torch_gpu,
    require_torch_or_tf,
)

from .test_pipelines_common import ANY


@is_pipeline_test
@require_torch_or_tf
class 			SCREAMING_SNAKE_CASE      (      unittest.TestCase       ):
     '''simple docstring'''


     UpperCamelCase_				:		Union[str, Any]								=						MODEL_FOR_CAUSAL_LM_MAPPING
     UpperCamelCase_				:		Optional[int]								=						TF_MODEL_FOR_CAUSAL_LM_MAPPING


     @require_torch
     def 					_A							(					self							:						str   ):
       SCREAMING_SNAKE_CASE							:   Optional[int]						    =     pipeline(task="text-generation"       ,			model="sshleifer/tiny-ctrl"       ,			framework="pt"   )
       # Using `do_sample=False` to force deterministic output
       SCREAMING_SNAKE_CASE							:   List[Any]						    =     text_generator("This is a test"       ,			do_sample=UpperCAmelCase_   )
       self.assertEqual(
           UpperCAmelCase_       ,			[
               {
                   "generated_text": (
                       "This is a test ☃ ☃ segmental segmental segmental 议议eski eski flutter flutter Lacy oscope."
                       " oscope. FiliFili@@"
                   )
               }
           ]       ,			)

       SCREAMING_SNAKE_CASE							:   str						    =     text_generator(["This is a test", "This is a second test"]   )
       self.assertEqual(
           UpperCAmelCase_       ,			[
               [
                   {
                       "generated_text": (
                           "This is a test ☃ ☃ segmental segmental segmental 议议eski eski flutter flutter Lacy oscope."
                           " oscope. FiliFili@@"
                       )
                   }
               ],
               [
                   {
                       "generated_text": (
                           "This is a second test ☃ segmental segmental segmental 议议eski eski flutter flutter Lacy"
                           " oscope. oscope. FiliFili@@"
                       )
                   }
               ],
           ]       ,			)

       SCREAMING_SNAKE_CASE							:   str						    =     text_generator("This is a test"       ,			do_sample=UpperCAmelCase_       ,			num_return_sequences=2       ,			return_tensors=UpperCAmelCase_   )
       self.assertEqual(
           UpperCAmelCase_       ,			[
               {"generated_token_ids": ANY(UpperCAmelCase_   )},
               {"generated_token_ids": ANY(UpperCAmelCase_   )},
           ]       ,			)
       SCREAMING_SNAKE_CASE							:   List[Any]						    =     text_generator.model.config.eos_token_id
       SCREAMING_SNAKE_CASE							:   Union[str, Any]						    =     "<pad>"
       SCREAMING_SNAKE_CASE							:   List[str]						    =     text_generator(
           ["This is a test", "This is a second test"]       ,			do_sample=UpperCAmelCase_       ,			num_return_sequences=2       ,			batch_size=2       ,			return_tensors=UpperCAmelCase_       ,			)
       self.assertEqual(
           UpperCAmelCase_       ,			[
               [
                   {"generated_token_ids": ANY(UpperCAmelCase_   )},
                   {"generated_token_ids": ANY(UpperCAmelCase_   )},
               ],
               [
                   {"generated_token_ids": ANY(UpperCAmelCase_   )},
                   {"generated_token_ids": ANY(UpperCAmelCase_   )},
               ],
           ]       ,			)


     @require_tf
     def 					_A							(					self							:						List[str]   ):
       SCREAMING_SNAKE_CASE							:   str						    =     pipeline(task="text-generation"       ,			model="sshleifer/tiny-ctrl"       ,			framework="tf"   )

       # Using `do_sample=False` to force deterministic output
       SCREAMING_SNAKE_CASE							:   str						    =     text_generator("This is a test"       ,			do_sample=UpperCAmelCase_   )
       self.assertEqual(
           UpperCAmelCase_       ,			[
               {
                   "generated_text": (
                       "This is a test FeyFeyFey(Croatis.), s.), Cannes Cannes Cannes 閲閲Cannes Cannes Cannes 攵"
                       " please,"
                   )
               }
           ]       ,			)

       SCREAMING_SNAKE_CASE							:   Optional[int]						    =     text_generator(["This is a test", "This is a second test"]       ,			do_sample=UpperCAmelCase_   )
       self.assertEqual(
           UpperCAmelCase_       ,			[
               [
                   {
                       "generated_text": (
                           "This is a test FeyFeyFey(Croatis.), s.), Cannes Cannes Cannes 閲閲Cannes Cannes Cannes 攵"
                           " please,"
                       )
                   }
               ],
               [
                   {
                       "generated_text": (
                           "This is a second test Chieftain Chieftain prefecture prefecture prefecture Cannes Cannes"
                           " Cannes 閲閲Cannes Cannes Cannes 攵 please,"
                       )
                   }
               ],
           ]       ,			)


     def 					_A							(					self							:						str       ,			UpperCAmelCase_							:						Optional[int]       ,			UpperCAmelCase_							:						List[str]       ,			UpperCAmelCase_							:						Any   ):
       SCREAMING_SNAKE_CASE							:   List[str]						    =     TextGenerationPipeline(model=UpperCAmelCase_       ,			tokenizer=UpperCAmelCase_   )
       return text_generator, ["This is a test", "Another test"]


     def 					_A							(					self							:						Union[str, Any]   ):
       SCREAMING_SNAKE_CASE							:   List[str]						    =     "Hello I believe in"
       SCREAMING_SNAKE_CASE							:   int						    =     pipeline("text-generation"       ,			model="hf-internal-testing/tiny-random-gpt2"   )
       SCREAMING_SNAKE_CASE							:   Optional[Any]						    =     text_generator(UpperCAmelCase_   )
       self.assertEqual(
           UpperCAmelCase_       ,			[{"generated_text": "Hello I believe in fe fe fe fe fe fe fe fe fe fe fe fe"}]       ,			)

       SCREAMING_SNAKE_CASE							:   Dict						    =     text_generator(UpperCAmelCase_       ,			stop_sequence=" fe"   )
       self.assertEqual(UpperCAmelCase_       ,			[{"generated_text": "Hello I believe in fe"}]   )


     def 					_A							(					self							:						int       ,			UpperCAmelCase_							:						int       ,			UpperCAmelCase_							:						Optional[int]   ):
       SCREAMING_SNAKE_CASE							:   Optional[Any]						    =     text_generator.model
       SCREAMING_SNAKE_CASE							:   Union[str, Any]						    =     text_generator.tokenizer

       SCREAMING_SNAKE_CASE							:   str						    =     text_generator("This is a test"   )
       self.assertEqual(UpperCAmelCase_       ,			[{"generated_text": ANY(UpperCAmelCase_   )}]   )
       self.assertTrue(outputs[0]["generated_text"].startswith("This is a test"   )   )

       SCREAMING_SNAKE_CASE							:   Optional[int]						    =     text_generator("This is a test"       ,			return_full_text=UpperCAmelCase_   )
       self.assertEqual(UpperCAmelCase_       ,			[{"generated_text": ANY(UpperCAmelCase_   )}]   )
       self.assertNotIn("This is a test"       ,			outputs[0]["generated_text"]   )

       SCREAMING_SNAKE_CASE							:   Optional[int]						    =     pipeline(task="text-generation"       ,			model=UpperCAmelCase_       ,			tokenizer=UpperCAmelCase_       ,			return_full_text=UpperCAmelCase_   )
       SCREAMING_SNAKE_CASE							:   Dict						    =     text_generator("This is a test"   )
       self.assertEqual(UpperCAmelCase_       ,			[{"generated_text": ANY(UpperCAmelCase_   )}]   )
       self.assertNotIn("This is a test"       ,			outputs[0]["generated_text"]   )

       SCREAMING_SNAKE_CASE							:   List[str]						    =     text_generator("This is a test"       ,			return_full_text=UpperCAmelCase_   )
       self.assertEqual(UpperCAmelCase_       ,			[{"generated_text": ANY(UpperCAmelCase_   )}]   )
       self.assertTrue(outputs[0]["generated_text"].startswith("This is a test"   )   )

       SCREAMING_SNAKE_CASE							:   int						    =     text_generator(["This is great !", "Something else"]       ,			num_return_sequences=2       ,			do_sample=UpperCAmelCase_   )
       self.assertEqual(
           UpperCAmelCase_       ,			[
               [{"generated_text": ANY(UpperCAmelCase_   )}, {"generated_text": ANY(UpperCAmelCase_   )}],
               [{"generated_text": ANY(UpperCAmelCase_   )}, {"generated_text": ANY(UpperCAmelCase_   )}],
           ]       ,			)

       if text_generator.tokenizer.pad_token is not None:
         SCREAMING_SNAKE_CASE							:   List[Any]						    =     text_generator(
             ["This is great !", "Something else"]       ,			num_return_sequences=2       ,			batch_size=2       ,			do_sample=UpperCAmelCase_   )
         self.assertEqual(
             UpperCAmelCase_       ,			[
                 [{"generated_text": ANY(UpperCAmelCase_   )}, {"generated_text": ANY(UpperCAmelCase_   )}],
                 [{"generated_text": ANY(UpperCAmelCase_   )}, {"generated_text": ANY(UpperCAmelCase_   )}],
             ]       ,			)

       with self.assertRaises(UpperCAmelCase_   ):
         SCREAMING_SNAKE_CASE							:   Dict						    =     text_generator("test"       ,			return_full_text=UpperCAmelCase_       ,			return_text=UpperCAmelCase_   )
       with self.assertRaises(UpperCAmelCase_   ):
         SCREAMING_SNAKE_CASE							:   Optional[Any]						    =     text_generator("test"       ,			return_full_text=UpperCAmelCase_       ,			return_tensors=UpperCAmelCase_   )
       with self.assertRaises(UpperCAmelCase_   ):
         SCREAMING_SNAKE_CASE							:   Optional[Any]						    =     text_generator("test"       ,			return_text=UpperCAmelCase_       ,			return_tensors=UpperCAmelCase_   )

       # Empty prompt is slighly special
       # it requires BOS token to exist.
       # Special case for Pegasus which will always append EOS so will
       # work even without BOS.
       if (
           text_generator.tokenizer.bos_token_id is not None
           or "Pegasus" in tokenizer.__class__.__name__
           or "Git" in model.__class__.__name__
       ):
         SCREAMING_SNAKE_CASE							:   Tuple						    =     text_generator(""   )
         self.assertEqual(UpperCAmelCase_       ,			[{"generated_text": ANY(UpperCAmelCase_   )}]   )
       else:
         with self.assertRaises((ValueError, AssertionError)   ):
           SCREAMING_SNAKE_CASE							:   Union[str, Any]						    =     text_generator(""   )

       if text_generator.framework == "tf":
         # TF generation does not support max_new_tokens, and it's impossible
         # to control long generation with only max_length without
         # fancy calculation, dismissing tests for now.
         return
       # We don't care about infinite range models.
       # They already work.
       # Skip this test for XGLM, since it uses sinusoidal positional embeddings which are resized on-the-fly.
       SCREAMING_SNAKE_CASE							:   Tuple						    =     ["RwkvForCausalLM", "XGLMForCausalLM", "GPTNeoXForCausalLM"]
       if (
           tokenizer.model_max_length < 1_0000
           and text_generator.model.__class__.__name__ not in EXTRA_MODELS_CAN_HANDLE_LONG_INPUTS
       ):
         # Handling of large generations
         with self.assertRaises((RuntimeError, IndexError, ValueError, AssertionError)   ):
           text_generator("This is a test" * 500       ,			max_new_tokens=20   )

         SCREAMING_SNAKE_CASE							:   Optional[Any]						    =     text_generator("This is a test" * 500       ,			handle_long_generation="hole"       ,			max_new_tokens=20   )
         # Hole strategy cannot work
         with self.assertRaises(UpperCAmelCase_   ):
           text_generator(
               "This is a test" * 500       ,			handle_long_generation="hole"       ,			max_new_tokens=tokenizer.model_max_length + 10       ,			)


     @require_torch
     @require_accelerate
     @require_torch_gpu
     def 					_A							(					self							:						Union[str, Any]   ):
       import torch

       # Classic `model_kwargs`
       SCREAMING_SNAKE_CASE							:   List[str]						    =     pipeline(
           model="hf-internal-testing/tiny-random-bloom"       ,			model_kwargs={"device_map": "auto", "torch_dtype": torch.bfloataa}       ,			)
       self.assertEqual(pipe.model.device       ,			torch.device(0   )   )
       self.assertEqual(pipe.model.lm_head.weight.dtype       ,			torch.bfloataa   )
       SCREAMING_SNAKE_CASE							:   List[str]						    =     pipe("This is a test"   )
       self.assertEqual(
           UpperCAmelCase_       ,			[
               {
                   "generated_text": (
                       "This is a test test test test test test test test test test test test test test test test"
                       " test"
                   )
               }
           ]       ,			)

       # Upgraded those two to real pipeline arguments (they just get sent for the model as they're unlikely to mean anything else.)
       SCREAMING_SNAKE_CASE							:   Tuple						    =     pipeline(model="hf-internal-testing/tiny-random-bloom"       ,			device_map="auto"       ,			torch_dtype=torch.bfloataa   )
       self.assertEqual(pipe.model.device       ,			torch.device(0   )   )
       self.assertEqual(pipe.model.lm_head.weight.dtype       ,			torch.bfloataa   )
       SCREAMING_SNAKE_CASE							:   str						    =     pipe("This is a test"   )
       self.assertEqual(
           UpperCAmelCase_       ,			[
               {
                   "generated_text": (
                       "This is a test test test test test test test test test test test test test test test test"
                       " test"
                   )
               }
           ]       ,			)

       # torch_dtype will be automatically set to float32 if not provided - check: https://github.com/huggingface/transformers/pull/20602
       SCREAMING_SNAKE_CASE							:   int						    =     pipeline(model="hf-internal-testing/tiny-random-bloom"       ,			device_map="auto"   )
       self.assertEqual(pipe.model.device       ,			torch.device(0   )   )
       self.assertEqual(pipe.model.lm_head.weight.dtype       ,			torch.floataa   )
       SCREAMING_SNAKE_CASE							:   List[Any]						    =     pipe("This is a test"   )
       self.assertEqual(
           UpperCAmelCase_       ,			[
               {
                   "generated_text": (
                       "This is a test test test test test test test test test test test test test test test test"
                       " test"
                   )
               }
           ]       ,			)


     @require_torch
     @require_torch_gpu
     def 					_A							(					self							:						Optional[Any]   ):
       import torch

       SCREAMING_SNAKE_CASE							:   Dict						    =     pipeline(model="hf-internal-testing/tiny-random-bloom"       ,			device=0       ,			torch_dtype=torch.floataa   )
       pipe("This is a test"   )


     @require_torch
     @require_accelerate
     @require_torch_gpu
     def 					_A							(					self							:						List[str]   ):
       import torch

       SCREAMING_SNAKE_CASE							:   Union[str, Any]						    =     pipeline(model="hf-internal-testing/tiny-random-bloom"       ,			device_map="auto"       ,			torch_dtype=torch.floataa   )
       pipe("This is a test"       ,			do_sample=UpperCAmelCase_       ,			top_p=0.5   )






     def 					_A							(					self							:						Any   ):
       SCREAMING_SNAKE_CASE							:   Optional[Any]						    =     "Hello world"
       SCREAMING_SNAKE_CASE							:   List[Any]						    =     pipeline("text-generation"       ,			model="hf-internal-testing/tiny-random-gpt2"   )
       if text_generator.model.framework == "tf":
         SCREAMING_SNAKE_CASE							:   str						    =     logging.get_logger("transformers.generation.tf_utils"   )
       else:
         SCREAMING_SNAKE_CASE							:   int						    =     logging.get_logger("transformers.generation.utils"   )
       SCREAMING_SNAKE_CASE							:   List[str]						    =     "Both `max_new_tokens`"  # The beggining of the message to be checked in this test

       # Both are set by the user -> log warning
       with CaptureLogger(UpperCAmelCase_   ) as cl:
         SCREAMING_SNAKE_CASE							:   int						    =     text_generator(UpperCAmelCase_       ,			max_length=10       ,			max_new_tokens=1   )
       self.assertIn(UpperCAmelCase_       ,			cl.out   )

       # The user only sets one -> no warning
       with CaptureLogger(UpperCAmelCase_   ) as cl:
         SCREAMING_SNAKE_CASE							:   Optional[Any]						    =     text_generator(UpperCAmelCase_       ,			max_new_tokens=1   )
       self.assertNotIn(UpperCAmelCase_       ,			cl.out   )

       with CaptureLogger(UpperCAmelCase_   ) as cl:
         SCREAMING_SNAKE_CASE							:   Optional[Any]						    =     text_generator(UpperCAmelCase_       ,			max_length=10   )
       self.assertNotIn(UpperCAmelCase_       ,			cl.out   )










'''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__       (			SCREAMING_SNAKE_CASE_						):



     """simple docstring"""






     __magic_name__           =     42
class 						UpperCamelCase__       (			SCREAMING_SNAKE_CASE_     ,       SCREAMING_SNAKE_CASE_						):



     """simple docstring"""






     @register_to_config
     def __init__(			self   ,						snake_case__ = 6_5536   ,						snake_case__ = None   ,						snake_case__ = 2   ,						snake_case__ = 2   ,						snake_case__ = 0   ,						snake_case__ = "fourier"   ,						snake_case__ = True   ,						snake_case__ = False   ,						snake_case__ = 0.0   ,						snake_case__ = ("DownBlock1DNoSkip", "DownBlock1D", "AttnDownBlock1D")   ,						snake_case__ = ("AttnUpBlock1D", "UpBlock1D", "UpBlock1DNoSkip")   ,						snake_case__ = "UNetMidBlock1D"   ,						snake_case__ = None   ,						snake_case__ = (32, 32, 64)   ,						snake_case__ = None   ,						snake_case__ = 8   ,						snake_case__ = 1   ,						snake_case__ = False   ,						):

            '''simple docstring'''

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

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

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

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

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

                   if i == 0:
                          input_channel += extra_in_channels

                   _lowerCAmelCase     :				Tuple       =							i == len(snake_case__ ) - 1

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

            # mid
            _lowerCAmelCase     :				List[str]       =							get_mid_block(
                snake_case__   ,						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=snake_case__   ,						add_downsample=snake_case__   ,						)

            # up
            _lowerCAmelCase     :				Union[str, Any]       =							list(reversed(snake_case__ ) )
            _lowerCAmelCase     :				List[Any]       =							reversed_block_out_channels[0]
            if out_block_type is None:
                   _lowerCAmelCase     :				Tuple       =							out_channels
            else:
                   _lowerCAmelCase     :				Tuple       =							block_out_channels[0]

            for i, up_block_type in enumerate(snake_case__ ):
                   _lowerCAmelCase     :				int       =							output_channel
                   _lowerCAmelCase     :				Dict       =							(
                       reversed_block_out_channels[i + 1] if i < len(snake_case__ ) - 1 else final_upsample_channels
                   )

                   _lowerCAmelCase     :				Union[str, Any]       =							i == len(snake_case__ ) - 1

                   _lowerCAmelCase     :				Tuple       =							get_up_block(
                       snake_case__   ,						num_layers=snake_case__   ,						in_channels=snake_case__   ,						out_channels=snake_case__   ,						temb_channels=block_out_channels[0]   ,						add_upsample=not is_final_block   ,						)
                   self.up_blocks.append(snake_case__ )
                   _lowerCAmelCase     :				List[str]       =							output_channel

            # out
            _lowerCAmelCase     :				Tuple       =							norm_num_groups if norm_num_groups is not None else min(block_out_channels[0] // 4   ,						32 )
            _lowerCAmelCase     :				Any       =							get_out_block(
                out_block_type=snake_case__   ,						num_groups_out=snake_case__   ,						embed_dim=block_out_channels[0]   ,						out_channels=snake_case__   ,						act_fn=snake_case__   ,						fc_dim=block_out_channels[-1] // 4   ,						)







     def a						(			self   ,						snake_case__   ,						snake_case__   ,						snake_case__ = True   ,						):

            '''simple docstring'''

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

            _lowerCAmelCase     :				int       =							self.time_proj(snake_case__ )
            if self.config.use_timestep_embedding:
                   _lowerCAmelCase     :				Optional[int]       =							self.time_mlp(snake_case__ )
            else:
                   _lowerCAmelCase     :				List[str]       =							timestep_embed[..., None]
                   _lowerCAmelCase     :				Dict       =							timestep_embed.repeat([1, 1, sample.shape[2]] ).to(sample.dtype )
                   _lowerCAmelCase     :				Tuple       =							timestep_embed.broadcast_to((sample.shape[:1] + timestep_embed.shape[1:]) )

            # 2. down
            _lowerCAmelCase     :				str       =							()
            for downsample_block in self.down_blocks:
                   _lowerCAmelCase      ,  _lowerCAmelCase     :				Union[str, Any]       =							downsample_block(hidden_states=snake_case__   ,						temb=snake_case__ )
                   down_block_res_samples += res_samples

            # 3. mid
            if self.mid_block:
                   _lowerCAmelCase     :				Any       =							self.mid_block(snake_case__   ,						snake_case__ )

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

            # 5. post-process
            if self.out_block:
                   _lowerCAmelCase     :				Optional[Any]       =							self.out_block(snake_case__   ,						snake_case__ )

            if not return_dict:
                   return (sample,)

            return UNetaDOutput(sample=snake_case__ )

