




import unittest

import numpy as np

from transformers import is_flax_available
from transformers.testing_utils import require_flax

from ..test_modeling_flax_common import ids_tensor


if is_flax_available():
     import jax
     import jax.numpy as jnp

     from transformers.generation import (
         FlaxForcedBOSTokenLogitsProcessor,
         FlaxForcedEOSTokenLogitsProcessor,
         FlaxLogitsProcessorList,
         FlaxMinLengthLogitsProcessor,
         FlaxTemperatureLogitsWarper,
         FlaxTopKLogitsWarper,
         FlaxTopPLogitsWarper,
     )






@require_flax
class  _lowerCAmelCase     (  unittest.TestCase	):

      def 	__magic_name__(   self      ,				__UpperCAmelCase      ,				__UpperCAmelCase		):
        lowerCAmelCase__				:  Tuple		       =     jnp.ones((batch_size, length)		) / length
        return scores

      def 	__magic_name__(   self		):
        lowerCAmelCase__				:  List[Any]		       =     None
        lowerCAmelCase__				:  Dict		       =     20

        lowerCAmelCase__				:  List[Any]		       =     self._get_uniform_logits(batch_size=2      ,				length=__UpperCAmelCase		)

        # tweak scores to not be uniform anymore
        lowerCAmelCase__				:  Optional[Any]		       =     scores.at[1, 5].set((1 / length) + 0.1		)  # peak, 1st batch
        lowerCAmelCase__				:  List[Any]		       =     scores.at[1, 10].set((1 / length) - 0.4		)  # valley, 1st batch

        # compute softmax
        lowerCAmelCase__				:  List[str]		       =     jax.nn.softmax(__UpperCAmelCase      ,				axis=-1		)

        lowerCAmelCase__				:  Union[str, Any]		       =     FlaxTemperatureLogitsWarper(temperature=0.5		)
        lowerCAmelCase__				:  List[str]		       =     FlaxTemperatureLogitsWarper(temperature=1.3		)

        lowerCAmelCase__				:  Union[str, Any]		       =     jax.nn.softmax(temp_dist_warper_sharper(__UpperCAmelCase      ,				scores.copy()      ,				cur_len=__UpperCAmelCase		)      ,				axis=-1		)
        lowerCAmelCase__				:  Tuple		       =     jax.nn.softmax(temp_dist_warper_smoother(__UpperCAmelCase      ,				scores.copy()      ,				cur_len=__UpperCAmelCase		)      ,				axis=-1		)

        # uniform distribution stays uniform
        self.assertTrue(jnp.allclose(probs[0, :]      ,				warped_prob_sharp[0, :]      ,				atol=1e-3		)		)
        self.assertTrue(jnp.allclose(probs[0, :]      ,				warped_prob_smooth[0, :]      ,				atol=1e-3		)		)

        # sharp peaks get higher, valleys get lower
        self.assertLess(probs[1, :].max()      ,				warped_prob_sharp[1, :].max()		)
        self.assertGreater(probs[1, :].min()      ,				warped_prob_sharp[1, :].min()		)

        # smooth peaks get lower, valleys get higher
        self.assertGreater(probs[1, :].max()      ,				warped_prob_smooth[1, :].max()		)
        self.assertLess(probs[1, :].min()      ,				warped_prob_smooth[1, :].min()		)

      def 	__magic_name__(   self		):
        lowerCAmelCase__				:  Union[str, Any]		       =     None
        lowerCAmelCase__				:  List[Any]		       =     10
        lowerCAmelCase__				:  Optional[Any]		       =     2

        # create ramp distribution
        lowerCAmelCase__				:  Any		       =     np.broadcast_to(np.arange(__UpperCAmelCase		)[None, :]      ,				(batch_size, vocab_size)		).copy()
        lowerCAmelCase__				:  Union[str, Any]		       =     ramp_logits[1:, : vocab_size // 2] + vocab_size

        lowerCAmelCase__				:  str		       =     FlaxTopKLogitsWarper(3		)

        lowerCAmelCase__				:  Optional[Any]		       =     top_k_warp(__UpperCAmelCase      ,				__UpperCAmelCase      ,				cur_len=__UpperCAmelCase		)

        # check that correct tokens are filtered
        self.assertListEqual(jnp.isinf(scores[0]		).tolist()      ,				7 * [True] + 3 * [False]		)
        self.assertListEqual(jnp.isinf(scores[1]		).tolist()      ,				2 * [True] + 3 * [False] + 5 * [True]		)

        # check special case
        lowerCAmelCase__				:  List[str]		       =     5
        lowerCAmelCase__				:  Union[str, Any]		       =     FlaxTopKLogitsWarper(top_k=1      ,				filter_value=0.0      ,				min_tokens_to_keep=3		)

        lowerCAmelCase__				:  Optional[Any]		       =     np.broadcast_to(np.arange(__UpperCAmelCase		)[None, :]      ,				(batch_size, length)		).copy()
        lowerCAmelCase__				:  int		       =     top_k_warp_safety_check(__UpperCAmelCase      ,				__UpperCAmelCase      ,				cur_len=__UpperCAmelCase		)

        # min_tokens overwrites k: 3 tokens are kept => 2 tokens are nullified
        self.assertListEqual((scores == 0.0).sum(axis=-1		).tolist()      ,				[2, 2]		)

      def 	__magic_name__(   self		):
        lowerCAmelCase__				:  Union[str, Any]		       =     None
        lowerCAmelCase__				:  Tuple		       =     10
        lowerCAmelCase__				:  Optional[Any]		       =     2

        # create distribution and take log (inverse to Softmax as taken in TopPLogitsWarper)
        lowerCAmelCase__				:  int		       =     np.log(np.array([[0.3, 0.1, 0.1, 0.5], [0.15, 0.3, 0.3, 0.25]]		)		)

        lowerCAmelCase__				:  str		       =     FlaxTopPLogitsWarper(0.8		)
        lowerCAmelCase__				:  Optional[int]		       =     np.exp(top_p_warp(__UpperCAmelCase      ,				__UpperCAmelCase      ,				cur_len=__UpperCAmelCase		)		)

        # dist should be filtered to keep min num values so that sum is >= top_p
        # exp (-inf) => 0
        lowerCAmelCase__				:  Any		       =     np.array([[0.3, 0.0, 0.0, 0.5], [0.0, 0.3, 0.3, 0.25]]		)
        self.assertTrue(np.allclose(__UpperCAmelCase      ,				__UpperCAmelCase      ,				atol=1e-3		)		)

        # check edge cases with negative and extreme logits
        lowerCAmelCase__				:  Union[str, Any]		       =     np.broadcast_to(np.arange(__UpperCAmelCase		)[None, :]      ,				(batch_size, vocab_size)		).copy() - (
            vocab_size // 2
        )

        # make ramp_logits more extreme
        lowerCAmelCase__				:  Optional[int]		       =     ramp_logits[1] * 100.0

        # make sure at least 2 tokens are kept
        lowerCAmelCase__				:  Tuple		       =     FlaxTopPLogitsWarper(0.9      ,				min_tokens_to_keep=2      ,				filter_value=0.0		)
        lowerCAmelCase__				:  List[Any]		       =     top_p_warp(__UpperCAmelCase      ,				__UpperCAmelCase      ,				cur_len=__UpperCAmelCase		)

        # first batch should keep three tokens, second batch would keep only 1, but due to `min_tokens_to_keep=2` keeps 2.
        self.assertListEqual((filtered_dist != 0.0).sum(axis=-1		).tolist()      ,				[3, 2]		)

      def 	__magic_name__(   self		):
        lowerCAmelCase__				:  Tuple		       =     20
        lowerCAmelCase__				:  Union[str, Any]		       =     4
        lowerCAmelCase__				:  str		       =     0

        lowerCAmelCase__				:  int		       =     FlaxMinLengthLogitsProcessor(min_length=10      ,				eos_token_id=__UpperCAmelCase		)

        # check that min length is applied at length 5
        lowerCAmelCase__				:  Tuple		       =     ids_tensor((batch_size, 20)      ,				vocab_size=20		)
        lowerCAmelCase__				:  int		       =     5
        lowerCAmelCase__				:  Any		       =     self._get_uniform_logits(__UpperCAmelCase      ,				__UpperCAmelCase		)
        lowerCAmelCase__				:  List[str]		       =     min_dist_processor(__UpperCAmelCase      ,				__UpperCAmelCase      ,				cur_len=__UpperCAmelCase		)
        self.assertListEqual(scores_before_min_length[:, eos_token_id].tolist()      ,				4 * [-float('''inf'''		)]		)

        # check that min length is not applied anymore at length 15
        lowerCAmelCase__				:  Dict		       =     self._get_uniform_logits(__UpperCAmelCase      ,				__UpperCAmelCase		)
        lowerCAmelCase__				:  Optional[int]		       =     15
        lowerCAmelCase__				:  Optional[Any]		       =     min_dist_processor(__UpperCAmelCase      ,				__UpperCAmelCase      ,				cur_len=__UpperCAmelCase		)
        self.assertFalse(jnp.isinf(__UpperCAmelCase		).any()		)

      def 	__magic_name__(   self		):
        lowerCAmelCase__				:  int		       =     20
        lowerCAmelCase__				:  int		       =     4
        lowerCAmelCase__				:  Union[str, Any]		       =     0

        lowerCAmelCase__				:  List[Any]		       =     FlaxForcedBOSTokenLogitsProcessor(bos_token_id=__UpperCAmelCase		)

        # check that all scores are -inf except the bos_token_id score
        lowerCAmelCase__				:  List[str]		       =     ids_tensor((batch_size, 1)      ,				vocab_size=20		)
        lowerCAmelCase__				:  Optional[Any]		       =     1
        lowerCAmelCase__				:  Tuple		       =     self._get_uniform_logits(__UpperCAmelCase      ,				__UpperCAmelCase		)
        lowerCAmelCase__				:  Union[str, Any]		       =     logits_processor(__UpperCAmelCase      ,				__UpperCAmelCase      ,				cur_len=__UpperCAmelCase		)
        self.assertTrue(jnp.isneginf(scores[:, bos_token_id + 1 :]		).all()		)
        self.assertListEqual(scores[:, bos_token_id].tolist()      ,				4 * [0]		)  # score for bos_token_id shold be zero

        # check that bos_token_id is not forced if current length is greater than 1
        lowerCAmelCase__				:  int		       =     3
        lowerCAmelCase__				:  str		       =     self._get_uniform_logits(__UpperCAmelCase      ,				__UpperCAmelCase		)
        lowerCAmelCase__				:  List[Any]		       =     logits_processor(__UpperCAmelCase      ,				__UpperCAmelCase      ,				cur_len=__UpperCAmelCase		)
        self.assertFalse(jnp.isinf(__UpperCAmelCase		).any()		)

      def 	__magic_name__(   self		):
        lowerCAmelCase__				:  int		       =     20
        lowerCAmelCase__				:  List[Any]		       =     4
        lowerCAmelCase__				:  Any		       =     0
        lowerCAmelCase__				:  List[str]		       =     5

        lowerCAmelCase__				:  Optional[Any]		       =     FlaxForcedEOSTokenLogitsProcessor(max_length=__UpperCAmelCase      ,				eos_token_id=__UpperCAmelCase		)

        # check that all scores are -inf except the eos_token_id when max_length is reached
        lowerCAmelCase__				:  Optional[int]		       =     ids_tensor((batch_size, 4)      ,				vocab_size=20		)
        lowerCAmelCase__				:  Optional[int]		       =     4
        lowerCAmelCase__				:  Optional[int]		       =     self._get_uniform_logits(__UpperCAmelCase      ,				__UpperCAmelCase		)
        lowerCAmelCase__				:  Dict		       =     logits_processor(__UpperCAmelCase      ,				__UpperCAmelCase      ,				cur_len=__UpperCAmelCase		)
        self.assertTrue(jnp.isneginf(scores[:, eos_token_id + 1 :]		).all()		)
        self.assertListEqual(scores[:, eos_token_id].tolist()      ,				4 * [0]		)  # score for eos_token_id should be zero

        # check that eos_token_id is not forced if max_length is not reached
        lowerCAmelCase__				:  Optional[Any]		       =     3
        lowerCAmelCase__				:  Union[str, Any]		       =     self._get_uniform_logits(__UpperCAmelCase      ,				__UpperCAmelCase		)
        lowerCAmelCase__				:  List[Any]		       =     logits_processor(__UpperCAmelCase      ,				__UpperCAmelCase      ,				cur_len=__UpperCAmelCase		)
        self.assertFalse(jnp.isinf(__UpperCAmelCase		).any()		)

      def 	__magic_name__(   self		):
        lowerCAmelCase__				:  Union[str, Any]		       =     4
        lowerCAmelCase__				:  Union[str, Any]		       =     10
        lowerCAmelCase__				:  List[Any]		       =     15
        lowerCAmelCase__				:  Dict		       =     2
        lowerCAmelCase__				:  List[Any]		       =     1
        lowerCAmelCase__				:  Any		       =     15

        # dummy input_ids and scores
        lowerCAmelCase__				:  int		       =     ids_tensor((batch_size, sequence_length)      ,				__UpperCAmelCase		)
        lowerCAmelCase__				:  Dict		       =     input_ids.copy()

        lowerCAmelCase__				:  List[str]		       =     self._get_uniform_logits(__UpperCAmelCase      ,				__UpperCAmelCase		)
        lowerCAmelCase__				:  Tuple		       =     scores.copy()

        # instantiate all dist processors
        lowerCAmelCase__				:  List[Any]		       =     FlaxTemperatureLogitsWarper(temperature=0.5		)
        lowerCAmelCase__				:  int		       =     FlaxTopKLogitsWarper(3		)
        lowerCAmelCase__				:  List[Any]		       =     FlaxTopPLogitsWarper(0.8		)

        # instantiate all logits processors
        lowerCAmelCase__				:  Optional[Any]		       =     FlaxMinLengthLogitsProcessor(min_length=10      ,				eos_token_id=__UpperCAmelCase		)
        lowerCAmelCase__				:  Optional[int]		       =     FlaxForcedBOSTokenLogitsProcessor(bos_token_id=__UpperCAmelCase		)
        lowerCAmelCase__				:  Optional[int]		       =     FlaxForcedEOSTokenLogitsProcessor(max_length=__UpperCAmelCase      ,				eos_token_id=__UpperCAmelCase		)

        lowerCAmelCase__				:  List[str]		       =     10

        # no processor list
        lowerCAmelCase__				:  int		       =     temp_dist_warp(__UpperCAmelCase      ,				__UpperCAmelCase      ,				cur_len=__UpperCAmelCase		)
        lowerCAmelCase__				:  Tuple		       =     top_k_warp(__UpperCAmelCase      ,				__UpperCAmelCase      ,				cur_len=__UpperCAmelCase		)
        lowerCAmelCase__				:  Optional[int]		       =     top_p_warp(__UpperCAmelCase      ,				__UpperCAmelCase      ,				cur_len=__UpperCAmelCase		)
        lowerCAmelCase__				:  Union[str, Any]		       =     min_dist_proc(__UpperCAmelCase      ,				__UpperCAmelCase      ,				cur_len=__UpperCAmelCase		)
        lowerCAmelCase__				:  Union[str, Any]		       =     bos_dist_proc(__UpperCAmelCase      ,				__UpperCAmelCase      ,				cur_len=__UpperCAmelCase		)
        lowerCAmelCase__				:  Any		       =     eos_dist_proc(__UpperCAmelCase      ,				__UpperCAmelCase      ,				cur_len=__UpperCAmelCase		)

        # with processor list
        lowerCAmelCase__				:  Union[str, Any]		       =     FlaxLogitsProcessorList(
            [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]		)
        lowerCAmelCase__				:  List[str]		       =     processor(__UpperCAmelCase      ,				__UpperCAmelCase      ,				cur_len=__UpperCAmelCase		)

        # scores should be equal
        self.assertTrue(jnp.allclose(__UpperCAmelCase      ,				__UpperCAmelCase      ,				atol=1e-3		)		)

        # input_ids should never be changed
        self.assertListEqual(input_ids.tolist()      ,				input_ids_comp.tolist()		)







      def 	__magic_name__(   self		):
        lowerCAmelCase__				:  List[Any]		       =     4
        lowerCAmelCase__				:  Dict		       =     10
        lowerCAmelCase__				:  Optional[int]		       =     15
        lowerCAmelCase__				:  List[str]		       =     2
        lowerCAmelCase__				:  Any		       =     1
        lowerCAmelCase__				:  List[Any]		       =     15

        # dummy input_ids and scores
        lowerCAmelCase__				:  List[str]		       =     ids_tensor((batch_size, sequence_length)      ,				__UpperCAmelCase		)
        lowerCAmelCase__				:  List[Any]		       =     input_ids.copy()

        lowerCAmelCase__				:  List[Any]		       =     self._get_uniform_logits(__UpperCAmelCase      ,				__UpperCAmelCase		)
        lowerCAmelCase__				:  int		       =     scores.copy()

        # instantiate all dist processors
        lowerCAmelCase__				:  List[Any]		       =     FlaxTemperatureLogitsWarper(temperature=0.5		)
        lowerCAmelCase__				:  List[str]		       =     FlaxTopKLogitsWarper(3		)
        lowerCAmelCase__				:  Optional[Any]		       =     FlaxTopPLogitsWarper(0.8		)

        # instantiate all logits processors
        lowerCAmelCase__				:  Optional[int]		       =     FlaxMinLengthLogitsProcessor(min_length=10      ,				eos_token_id=__UpperCAmelCase		)
        lowerCAmelCase__				:  Dict		       =     FlaxForcedBOSTokenLogitsProcessor(bos_token_id=__UpperCAmelCase		)
        lowerCAmelCase__				:  str		       =     FlaxForcedEOSTokenLogitsProcessor(max_length=__UpperCAmelCase      ,				eos_token_id=__UpperCAmelCase		)

        lowerCAmelCase__				:  Dict		       =     10

        # no processor list
        def run_no_processor_list(__UpperCAmelCase      ,				__UpperCAmelCase      ,				__UpperCAmelCase		):
          lowerCAmelCase__				:  Optional[int]		       =     temp_dist_warp(__UpperCAmelCase      ,				__UpperCAmelCase      ,				cur_len=__UpperCAmelCase		)
          lowerCAmelCase__				:  int		       =     top_k_warp(__UpperCAmelCase      ,				__UpperCAmelCase      ,				cur_len=__UpperCAmelCase		)
          lowerCAmelCase__				:  Tuple		       =     top_p_warp(__UpperCAmelCase      ,				__UpperCAmelCase      ,				cur_len=__UpperCAmelCase		)
          lowerCAmelCase__				:  List[Any]		       =     min_dist_proc(__UpperCAmelCase      ,				__UpperCAmelCase      ,				cur_len=__UpperCAmelCase		)
          lowerCAmelCase__				:  int		       =     bos_dist_proc(__UpperCAmelCase      ,				__UpperCAmelCase      ,				cur_len=__UpperCAmelCase		)
          lowerCAmelCase__				:  int		       =     eos_dist_proc(__UpperCAmelCase      ,				__UpperCAmelCase      ,				cur_len=__UpperCAmelCase		)
          return scores

        # with processor list
        def run_processor_list(__UpperCAmelCase      ,				__UpperCAmelCase      ,				__UpperCAmelCase		):
          lowerCAmelCase__				:  Optional[int]		       =     FlaxLogitsProcessorList(
              [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]		)
          lowerCAmelCase__				:  Tuple		       =     processor(__UpperCAmelCase      ,				__UpperCAmelCase      ,				cur_len=__UpperCAmelCase		)
          return scores

        lowerCAmelCase__				:  Union[str, Any]		       =     jax.jit(__UpperCAmelCase		)
        lowerCAmelCase__				:  List[str]		       =     jax.jit(__UpperCAmelCase		)

        lowerCAmelCase__				:  Dict		       =     jitted_run_no_processor_list(__UpperCAmelCase      ,				__UpperCAmelCase      ,				__UpperCAmelCase		)
        lowerCAmelCase__				:  Optional[int]		       =     jitted_run_processor_list(__UpperCAmelCase      ,				__UpperCAmelCase      ,				__UpperCAmelCase		)

        # scores should be equal
        self.assertTrue(jnp.allclose(__UpperCAmelCase      ,				__UpperCAmelCase      ,				atol=1e-3		)		)

        # input_ids should never be changed
        self.assertListEqual(input_ids.tolist()      ,				input_ids_comp.tolist()		)







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  _lowerCAmelCase     (  _lowercase	):
      A__		     =		42






class  _lowerCAmelCase     (  _lowercase    ,			_lowercase	):

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

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

        if use_timestep_embedding:
          lowerCAmelCase__				:  str		       =     block_out_channels[0] * 4
          lowerCAmelCase__				:  Dict		       =     TimestepEmbedding(
              in_channels=__UpperCAmelCase      ,				time_embed_dim=__UpperCAmelCase      ,				act_fn=__UpperCAmelCase      ,				out_dim=block_out_channels[0]      ,				)

        lowerCAmelCase__				:  str		       =     nn.ModuleList([]		)
        lowerCAmelCase__				:  Optional[Any]		       =     None
        lowerCAmelCase__				:  Optional[int]		       =     nn.ModuleList([]		)
        lowerCAmelCase__				:  Optional[Any]		       =     None

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

          if i == 0:
            input_channel += extra_in_channels

          lowerCAmelCase__				:  Any		       =     i == len(__UpperCAmelCase		) - 1

          lowerCAmelCase__				:  Any		       =     get_down_block(
              __UpperCAmelCase      ,				num_layers=__UpperCAmelCase      ,				in_channels=__UpperCAmelCase      ,				out_channels=__UpperCAmelCase      ,				temb_channels=block_out_channels[0]      ,				add_downsample=not is_final_block or downsample_each_block      ,				)
          self.down_blocks.append(__UpperCAmelCase		)

        # mid
        lowerCAmelCase__				:  Optional[int]		       =     get_mid_block(
            __UpperCAmelCase      ,				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=__UpperCAmelCase      ,				add_downsample=__UpperCAmelCase      ,				)

        # up
        lowerCAmelCase__				:  Union[str, Any]		       =     list(reversed(__UpperCAmelCase		)		)
        lowerCAmelCase__				:  List[str]		       =     reversed_block_out_channels[0]
        if out_block_type is None:
          lowerCAmelCase__				:  Optional[int]		       =     out_channels
        else:
          lowerCAmelCase__				:  int		       =     block_out_channels[0]

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

          lowerCAmelCase__				:  Optional[int]		       =     i == len(__UpperCAmelCase		) - 1

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

        # out
        lowerCAmelCase__				:  Union[str, Any]		       =     norm_num_groups if norm_num_groups is not None else min(block_out_channels[0] // 4      ,				32		)
        lowerCAmelCase__				:  Dict		       =     get_out_block(
            out_block_type=__UpperCAmelCase      ,				num_groups_out=__UpperCAmelCase      ,				embed_dim=block_out_channels[0]      ,				out_channels=__UpperCAmelCase      ,				act_fn=__UpperCAmelCase      ,				fc_dim=block_out_channels[-1] // 4      ,				)







      def 	__magic_name__(   self      ,				__UpperCAmelCase      ,				__UpperCAmelCase      ,				__UpperCAmelCase = True      ,				):
        lowerCAmelCase__				:  int		       =     timestep
        if not torch.is_tensor(__UpperCAmelCase		):
          lowerCAmelCase__				:  Any		       =     torch.tensor([timesteps]      ,				dtype=torch.long      ,				device=sample.device		)
        elif torch.is_tensor(__UpperCAmelCase		) and len(timesteps.shape		) == 0:
          lowerCAmelCase__				:  Any		       =     timesteps[None].to(sample.device		)

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

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

        # 3. mid
        if self.mid_block:
          lowerCAmelCase__				:  Dict		       =     self.mid_block(__UpperCAmelCase      ,				__UpperCAmelCase		)

        # 4. up
        for i, upsample_block in enumerate(self.up_blocks		):
          lowerCAmelCase__				:  Tuple		       =     down_block_res_samples[-1:]
          lowerCAmelCase__				:  int		       =     down_block_res_samples[:-1]
          lowerCAmelCase__				:  Optional[int]		       =     upsample_block(__UpperCAmelCase      ,				res_hidden_states_tuple=__UpperCAmelCase      ,				temb=__UpperCAmelCase		)

        # 5. post-process
        if self.out_block:
          lowerCAmelCase__				:  Any		       =     self.out_block(__UpperCAmelCase      ,				__UpperCAmelCase		)

        if not return_dict:
          return (sample,)

        return UNetaDOutput(sample=__UpperCAmelCase		)

