





import argparse
import os
import re


_UpperCAmelCase       : Any			     =       """src/transformers"""

# Pattern that looks at the indentation in a line.
_UpperCAmelCase       : Optional[int]			     =       re.compile(R"""^(\s*)\S""")
# Pattern that matches `"key":" and puts `key` in group 0.
_UpperCAmelCase       : Tuple			     =       re.compile(R"""^\s*\"([^\"]+)\":""")
# Pattern that matches `_import_structure["key"]` and puts `key` in group 0.
_UpperCAmelCase       : Any			     =       re.compile(R"""^\s*_import_structure\[\"([^\"]+)\"\]""")
# Pattern that matches `"key",` and puts `key` in group 0.
_UpperCAmelCase       : Tuple			     =       re.compile(R"""^\s*\"([^\"]+)\",\s*$""")
# Pattern that matches any `[stuff]` and puts `stuff` in group 0.
_UpperCAmelCase       : Any			     =       re.compile(R"""\[([^\]]+)\]""")


def  SCREAMING_SNAKE_CASE   ( _UpperCAmelCase				)    ->  int:
    lowerCamelCase__					:   Optional[Any]      				=	_re_indent.search(_UpperCAmelCase				)
    return "" if search is None else search.groups()[0]


def  SCREAMING_SNAKE_CASE   ( _UpperCAmelCase		,			_UpperCAmelCase=""		,			_UpperCAmelCase=None		,			_UpperCAmelCase=None				)    ->  Dict:
    lowerCamelCase__					:   int      				=	0
    lowerCamelCase__					:   int      				=	code.split('\n'				)
    if start_prompt is not None:
        while not lines[index].startswith(_UpperCAmelCase				):
            index += 1
        lowerCamelCase__					:   Union[str, Any]      				=	['\n'.join(lines[:index]				)]
    else:
        lowerCamelCase__					:   List[str]      				=	[]

    # We split into blocks until we get to the `end_prompt` (or the end of the block).
    lowerCamelCase__					:   Any      				=	[lines[index]]
    index += 1
    while index < len(_UpperCAmelCase				) and (end_prompt is None or not lines[index].startswith(_UpperCAmelCase				)):
        if len(lines[index]				) > 0 and get_indent(lines[index]				) == indent_level:
            if len(_UpperCAmelCase				) > 0 and get_indent(current_block[-1]				).startswith(indent_level + ' '				):
                current_block.append(lines[index]				)
                blocks.append('\n'.join(_UpperCAmelCase				)				)
                if index < len(_UpperCAmelCase				) - 1:
                    lowerCamelCase__					:   Optional[Any]      				=	[lines[index + 1]]
                    index += 1
                else:
                    lowerCamelCase__					:   List[Any]      				=	[]
            else:
                blocks.append('\n'.join(_UpperCAmelCase				)				)
                lowerCamelCase__					:   int      				=	[lines[index]]
        else:
            current_block.append(lines[index]				)
        index += 1

    # Adds current block if it's nonempty.
    if len(_UpperCAmelCase				) > 0:
        blocks.append('\n'.join(_UpperCAmelCase				)				)

    # Add final block after end_prompt if provided.
    if end_prompt is not None and index < len(_UpperCAmelCase				):
        blocks.append('\n'.join(lines[index:]				)				)

    return blocks


def  SCREAMING_SNAKE_CASE   ( _UpperCAmelCase				)    ->  int:

    def _inner(_UpperCAmelCase				):
        return key(_UpperCAmelCase				).lower().replace('_'		,			''				)

    return _inner


def  SCREAMING_SNAKE_CASE   ( _UpperCAmelCase		,			_UpperCAmelCase=None				)    ->  Union[str, Any]:

    # If no key is provided, we use a noop.
    def noop(_UpperCAmelCase				):
        return x

    if key is None:
        lowerCamelCase__					:   Optional[int]      				=	noop
    # Constants are all uppercase, they go first.
    lowerCamelCase__					:   List[str]      				=	[obj for obj in objects if key(_UpperCAmelCase				).isupper()]
    # Classes are not all uppercase but start with a capital, they go second.
    lowerCamelCase__					:   Any      				=	[obj for obj in objects if key(_UpperCAmelCase				)[0].isupper() and not key(_UpperCAmelCase				).isupper()]
    # Functions begin with a lowercase, they go last.
    lowerCamelCase__					:   Dict      				=	[obj for obj in objects if not key(_UpperCAmelCase				)[0].isupper()]

    lowerCamelCase__					:   str      				=	ignore_underscore(_UpperCAmelCase				)
    return sorted(_UpperCAmelCase		,			key=_UpperCAmelCase				) + sorted(_UpperCAmelCase		,			key=_UpperCAmelCase				) + sorted(_UpperCAmelCase		,			key=_UpperCAmelCase				)


def  SCREAMING_SNAKE_CASE   ( _UpperCAmelCase				)    ->  Union[str, Any]:

    # This inner function sort imports between [ ].
    def _replace(_UpperCAmelCase				):
        lowerCamelCase__					:   Union[str, Any]      				=	match.groups()[0]
        if "," not in imports:
            return F"""[{imports}]"""
        lowerCamelCase__					:   Optional[int]      				=	[part.strip().replace('"'		,			''				) for part in imports.split(','				)]
        # We will have a final empty element if the line finished with a comma.
        if len(keys[-1]				) == 0:
            lowerCamelCase__					:   Dict      				=	keys[:-1]
        return "[" + ", ".join([F"""\"{k}\"""" for k in sort_objects(_UpperCAmelCase				)]				) + "]"

    lowerCamelCase__					:   List[str]      				=	import_statement.split('\n'				)
    if len(_UpperCAmelCase				) > 3:
        # Here we have to sort internal imports that are on several lines (one per name):
        # key: [
        #     "object1",
        #     "object2",
        #     ...
        # ]

        # We may have to ignore one or two lines on each side.
        lowerCamelCase__					:   str      				=	2 if lines[1].strip() == '[' else 1
        lowerCamelCase__					:   int      				=	[(i, _re_strip_line.search(_UpperCAmelCase				).groups()[0]) for i, line in enumerate(lines[idx:-idx]				)]
        lowerCamelCase__					:   int      				=	sort_objects(_UpperCAmelCase		,			key=lambda _UpperCAmelCase				: x[1]				)
        lowerCamelCase__					:   str      				=	[lines[x[0] + idx] for x in sorted_indices]
        return "\n".join(lines[:idx] + sorted_lines + lines[-idx:]				)
    elif len(_UpperCAmelCase				) == 3:
        # Here we have to sort internal imports that are on one separate line:
        # key: [
        #     "object1", "object2", ...
        # ]
        if _re_bracket_content.search(lines[1]				) is not None:
            lowerCamelCase__					:   Dict      				=	_re_bracket_content.sub(_replace		,			lines[1]				)
        else:
            lowerCamelCase__					:   Dict      				=	[part.strip().replace('"'		,			''				) for part in lines[1].split(','				)]
            # We will have a final empty element if the line finished with a comma.
            if len(keys[-1]				) == 0:
                lowerCamelCase__					:   List[str]      				=	keys[:-1]
            lowerCamelCase__					:   Tuple      				=	get_indent(lines[1]				) + ', '.join([F"""\"{k}\"""" for k in sort_objects(_UpperCAmelCase				)]				)
        return "\n".join(_UpperCAmelCase				)
    else:
        # Finally we have to deal with imports fitting on one line
        lowerCamelCase__					:   Optional[Any]      				=	_re_bracket_content.sub(_replace		,			_UpperCAmelCase				)
        return import_statement


def  SCREAMING_SNAKE_CASE   ( _UpperCAmelCase		,			_UpperCAmelCase=True				)    ->  Tuple:
    with open(_UpperCAmelCase		,			encoding='utf-8'				) as f:
        lowerCamelCase__					:   Dict      				=	f.read()

    if "_import_structure" not in code:
        return

    # Blocks of indent level 0
    lowerCamelCase__					:   Optional[Any]      				=	split_code_in_indented_blocks(
        _UpperCAmelCase		,			start_prompt='_import_structure = {'		,			end_prompt='if TYPE_CHECKING:'				)

    # We ignore block 0 (everything untils start_prompt) and the last block (everything after end_prompt).
    for block_idx in range(1		,			len(_UpperCAmelCase				) - 1				):
        # Check if the block contains some `_import_structure`s thingy to sort.
        lowerCamelCase__					:   Dict      				=	main_blocks[block_idx]
        lowerCamelCase__					:   List[str]      				=	block.split('\n'				)

        # Get to the start of the imports.
        lowerCamelCase__					:   Union[str, Any]      				=	0
        while line_idx < len(_UpperCAmelCase				) and "_import_structure" not in block_lines[line_idx]:
            # Skip dummy import blocks
            if "import dummy" in block_lines[line_idx]:
                lowerCamelCase__					:   str      				=	len(_UpperCAmelCase				)
            else:
                line_idx += 1
        if line_idx >= len(_UpperCAmelCase				):
            continue

        # Ignore beginning and last line: they don't contain anything.
        lowerCamelCase__					:   Any      				=	'\n'.join(block_lines[line_idx:-1]				)
        lowerCamelCase__					:   int      				=	get_indent(block_lines[1]				)
        # Slit the internal block into blocks of indent level 1.
        lowerCamelCase__					:   Dict      				=	split_code_in_indented_blocks(_UpperCAmelCase		,			indent_level=_UpperCAmelCase				)
        # We have two categories of import key: list or _import_structure[key].append/extend
        lowerCamelCase__					:   List[str]      				=	_re_direct_key if '_import_structure = {' in block_lines[0] else _re_indirect_key
        # Grab the keys, but there is a trap: some lines are empty or just comments.
        lowerCamelCase__					:   Union[str, Any]      				=	[(pattern.search(_UpperCAmelCase				).groups()[0] if pattern.search(_UpperCAmelCase				) is not None else None) for b in internal_blocks]
        # We only sort the lines with a key.
        lowerCamelCase__					:   Optional[int]      				=	[(i, key) for i, key in enumerate(_UpperCAmelCase				) if key is not None]
        lowerCamelCase__					:   Optional[Any]      				=	[x[0] for x in sorted(_UpperCAmelCase		,			key=lambda _UpperCAmelCase				: x[1]				)]

        # We reorder the blocks by leaving empty lines/comments as they were and reorder the rest.
        lowerCamelCase__					:   str      				=	0
        lowerCamelCase__					:   Optional[Any]      				=	[]
        for i in range(len(_UpperCAmelCase				)				):
            if keys[i] is None:
                reorderded_blocks.append(internal_blocks[i]				)
            else:
                lowerCamelCase__					:   List[str]      				=	sort_objects_in_import(internal_blocks[sorted_indices[count]]				)
                reorderded_blocks.append(_UpperCAmelCase				)
                count += 1

        # And we put our main block back together with its first and last line.
        lowerCamelCase__					:   Dict      				=	'\n'.join(block_lines[:line_idx] + reorderded_blocks + [block_lines[-1]]				)

    if code != "\n".join(_UpperCAmelCase				):
        if check_only:
            return True
        else:
            print(F"""Overwriting {file}."""				)
            with open(_UpperCAmelCase		,			'w'		,			encoding='utf-8'				) as f:
                f.write('\n'.join(_UpperCAmelCase				)				)


def  SCREAMING_SNAKE_CASE   ( _UpperCAmelCase=True				)    ->  List[Any]:
    lowerCamelCase__					:   int      				=	[]
    for root, _, files in os.walk(_UpperCAmelCase				):
        if "__init__.py" in files:
            lowerCamelCase__					:   Optional[int]      				=	sort_imports(os.path.join(_UpperCAmelCase		,			'__init__.py'				)		,			check_only=_UpperCAmelCase				)
            if result:
                lowerCamelCase__					:   Optional[Any]      				=	[os.path.join(_UpperCAmelCase		,			'__init__.py'				)]
    if len(_UpperCAmelCase				) > 0:
        raise ValueError(F"""Would overwrite {len(_UpperCAmelCase				)} files, run `make style`."""				)


if __name__ == "__main__":
    _UpperCAmelCase       : Union[str, Any]			     =       argparse.ArgumentParser()
    parser.add_argument("""--check_only""", action="""store_true""", help="""Whether to only check or fix style.""")
    _UpperCAmelCase       : List[Any]			     =       parser.parse_args()

    sort_imports_in_all_inits(check_only=args.check_only)












from typing import Optional, Tuple, Union

import flax
import flax.linen as nn
import jax
import jax.numpy as jnp
from flax.core.frozen_dict import FrozenDict

from ..configuration_utils import ConfigMixin, flax_register_to_config
from ..utils import BaseOutput
from .embeddings_flax import FlaxTimestepEmbedding, FlaxTimesteps
from .modeling_flax_utils import FlaxModelMixin
from .unet_ad_blocks_flax import (
    FlaxCrossAttnDownBlockaD,
    FlaxCrossAttnUpBlockaD,
    FlaxDownBlockaD,
    FlaxUNetMidBlockaDCrossAttn,
    FlaxUpBlockaD,
)

@flax.struct.dataclass
class      lowerCAmelCase							(      __UpperCamelCase    ):
 UpperCAmelCase__									=	42


@flax_register_to_config
class      lowerCAmelCase							(      nn.Module,						__UpperCamelCase,						__UpperCamelCase    ):
 UpperCAmelCase__									=	32
 UpperCAmelCase__									=	4
 UpperCAmelCase__									=	4
 UpperCAmelCase__									=	(
     "CrossAttnDownBlock2D",
     "CrossAttnDownBlock2D",
     "CrossAttnDownBlock2D",
     "DownBlock2D",
 )
 UpperCAmelCase__									=	("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")
 UpperCAmelCase__									=	False
 UpperCAmelCase__									=	(3_20, 6_40, 12_80, 12_80)
 UpperCAmelCase__									=	2
 UpperCAmelCase__									=	8
 UpperCAmelCase__									=	None
 UpperCAmelCase__									=	12_80
 UpperCAmelCase__									=	0.0
 UpperCAmelCase__									=	False
 UpperCAmelCase__									=	jnp.floataa
 UpperCAmelCase__									=	True
 UpperCAmelCase__									=	0
 UpperCAmelCase__									=	False
 def   A_			(   self						:   Tuple    ,  UpperCAmelCase						:   jax.random.KeyArray			)    ->			FrozenDict:
     # init input tensors
     lowerCamelCase__					:   int      				=	(1, self.in_channels, self.sample_size, self.sample_size)
     lowerCamelCase__					:   List[str]      				=	jnp.zeros(UpperCAmelCase    ,  dtype=jnp.floataa			)
     lowerCamelCase__					:   Tuple      				=	jnp.ones((1,)    ,  dtype=jnp.intaa			)
     lowerCamelCase__					:   Dict      				=	jnp.zeros((1, 1, self.cross_attention_dim)    ,  dtype=jnp.floataa			)

     lowerCamelCase__				,							lowerCamelCase__					:   Union[str, Any]      				=	jax.random.split(UpperCAmelCase			)
     lowerCamelCase__					:   Dict      				=	{'params': params_rng, 'dropout': dropout_rng}

     return self.init(UpperCAmelCase    ,  UpperCAmelCase    ,  UpperCAmelCase    ,  UpperCAmelCase			)["params"]
 def   A_			(   self						:   Tuple			)    ->			Optional[int]:
     lowerCamelCase__					:   Any      				=	self.block_out_channels
     lowerCamelCase__					:   int      				=	block_out_channels[0] * 4

     if self.num_attention_heads is not None:
         raise ValueError(
             'At the moment it is not possible to define the number of attention heads via `num_attention_heads` because of a naming issue as described in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131. Passing `num_attention_heads` will only be supported in diffusers v0.19.'			)

     # If `num_attention_heads` is not defined (which is the case for most models)
     # it will default to `attention_head_dim`. This looks weird upon first reading it and it is.
     # The reason for this behavior is to correct for incorrectly named variables that were introduced
     # when this library was created. The incorrect naming was only discovered much later in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131
     # Changing `attention_head_dim` to `num_attention_heads` for 40,000+ configurations is too backwards breaking
     # which is why we correct for the naming here.
     lowerCamelCase__					:   Tuple      				=	self.num_attention_heads or self.attention_head_dim

     # input
     lowerCamelCase__					:   Optional[Any]      				=	nn.Conv(
         block_out_channels[0]    ,  kernel_size=(3, 3)    ,  strides=(1, 1)    ,  padding=((1, 1), (1, 1))    ,  dtype=self.dtype    ,  )

     # time
     lowerCamelCase__					:   Optional[int]      				=	FlaxTimesteps(
         block_out_channels[0]    ,  flip_sin_to_cos=self.flip_sin_to_cos    ,  freq_shift=self.config.freq_shift			)
     lowerCamelCase__					:   int      				=	FlaxTimestepEmbedding(UpperCAmelCase    ,  dtype=self.dtype			)

     lowerCamelCase__					:   Optional[int]      				=	self.only_cross_attention
     if isinstance(UpperCAmelCase    ,  UpperCAmelCase			):
         lowerCamelCase__					:   str      				=	(only_cross_attention,) * len(self.down_block_types			)

     if isinstance(UpperCAmelCase    ,  UpperCAmelCase			):
         lowerCamelCase__					:   List[Any]      				=	(num_attention_heads,) * len(self.down_block_types			)

     # down
     lowerCamelCase__					:   List[Any]      				=	[]
     lowerCamelCase__					:   Dict      				=	block_out_channels[0]
     for i, down_block_type in enumerate(self.down_block_types			):
         lowerCamelCase__					:   Dict      				=	output_channel
         lowerCamelCase__					:   Optional[int]      				=	block_out_channels[i]
         lowerCamelCase__					:   List[Any]      				=	i == len(UpperCAmelCase			) - 1

         if down_block_type == "CrossAttnDownBlock2D":
             lowerCamelCase__					:   Tuple      				=	FlaxCrossAttnDownBlockaD(
                 in_channels=UpperCAmelCase    ,  out_channels=UpperCAmelCase    ,  dropout=self.dropout    ,  num_layers=self.layers_per_block    ,  num_attention_heads=num_attention_heads[i]    ,  add_downsample=not is_final_block    ,  use_linear_projection=self.use_linear_projection    ,  only_cross_attention=only_cross_attention[i]    ,  use_memory_efficient_attention=self.use_memory_efficient_attention    ,  dtype=self.dtype    ,  )
         else:
             lowerCamelCase__					:   str      				=	FlaxDownBlockaD(
                 in_channels=UpperCAmelCase    ,  out_channels=UpperCAmelCase    ,  dropout=self.dropout    ,  num_layers=self.layers_per_block    ,  add_downsample=not is_final_block    ,  dtype=self.dtype    ,  )

         down_blocks.append(UpperCAmelCase			)
     lowerCamelCase__					:   List[Any]      				=	down_blocks

     # mid
     lowerCamelCase__					:   Dict      				=	FlaxUNetMidBlockaDCrossAttn(
         in_channels=block_out_channels[-1]    ,  dropout=self.dropout    ,  num_attention_heads=num_attention_heads[-1]    ,  use_linear_projection=self.use_linear_projection    ,  use_memory_efficient_attention=self.use_memory_efficient_attention    ,  dtype=self.dtype    ,  )

     # up
     lowerCamelCase__					:   Any      				=	[]
     lowerCamelCase__					:   Optional[int]      				=	list(reversed(UpperCAmelCase			)			)
     lowerCamelCase__					:   Any      				=	list(reversed(UpperCAmelCase			)			)
     lowerCamelCase__					:   int      				=	list(reversed(UpperCAmelCase			)			)
     lowerCamelCase__					:   Tuple      				=	reversed_block_out_channels[0]
     for i, up_block_type in enumerate(self.up_block_types			):
         lowerCamelCase__					:   str      				=	output_channel
         lowerCamelCase__					:   int      				=	reversed_block_out_channels[i]
         lowerCamelCase__					:   int      				=	reversed_block_out_channels[min(i + 1    ,  len(UpperCAmelCase			) - 1			)]

         lowerCamelCase__					:   Optional[Any]      				=	i == len(UpperCAmelCase			) - 1

         if up_block_type == "CrossAttnUpBlock2D":
             lowerCamelCase__					:   Tuple      				=	FlaxCrossAttnUpBlockaD(
                 in_channels=UpperCAmelCase    ,  out_channels=UpperCAmelCase    ,  prev_output_channel=UpperCAmelCase    ,  num_layers=self.layers_per_block + 1    ,  num_attention_heads=reversed_num_attention_heads[i]    ,  add_upsample=not is_final_block    ,  dropout=self.dropout    ,  use_linear_projection=self.use_linear_projection    ,  only_cross_attention=only_cross_attention[i]    ,  use_memory_efficient_attention=self.use_memory_efficient_attention    ,  dtype=self.dtype    ,  )
         else:
             lowerCamelCase__					:   Optional[Any]      				=	FlaxUpBlockaD(
                 in_channels=UpperCAmelCase    ,  out_channels=UpperCAmelCase    ,  prev_output_channel=UpperCAmelCase    ,  num_layers=self.layers_per_block + 1    ,  add_upsample=not is_final_block    ,  dropout=self.dropout    ,  dtype=self.dtype    ,  )

         up_blocks.append(UpperCAmelCase			)
         lowerCamelCase__					:   Tuple      				=	output_channel
     lowerCamelCase__					:   Tuple      				=	up_blocks

     # out
     lowerCamelCase__					:   Optional[int]      				=	nn.GroupNorm(num_groups=32    ,  epsilon=1e-5			)
     lowerCamelCase__					:   Any      				=	nn.Conv(
         self.out_channels    ,  kernel_size=(3, 3)    ,  strides=(1, 1)    ,  padding=((1, 1), (1, 1))    ,  dtype=self.dtype    ,  )







 def __call__(   self						:   int    ,  UpperCAmelCase						:   Any    ,  UpperCAmelCase						:   Tuple    ,  UpperCAmelCase						:   Optional[Any]    ,  UpperCAmelCase						:   Optional[Any]=None    ,  UpperCAmelCase						:   Dict=None    ,  UpperCAmelCase						:   bool = True    ,  UpperCAmelCase						:   bool = False    ,  )    ->			Union[FlaxUNetaDConditionOutput, Tuple]:
     # 1. time
     if not isinstance(UpperCAmelCase    ,  jnp.ndarray			):
         lowerCamelCase__					:   List[str]      				=	jnp.array([timesteps]    ,  dtype=jnp.intaa			)
     elif isinstance(UpperCAmelCase    ,  jnp.ndarray			) and len(timesteps.shape			) == 0:
         lowerCamelCase__					:   List[Any]      				=	timesteps.astype(dtype=jnp.floataa			)
         lowerCamelCase__					:   Any      				=	jnp.expand_dims(UpperCAmelCase    ,  0			)

     lowerCamelCase__					:   List[str]      				=	self.time_proj(UpperCAmelCase			)
     lowerCamelCase__					:   Optional[Any]      				=	self.time_embedding(UpperCAmelCase			)

     # 2. pre-process
     lowerCamelCase__					:   Dict      				=	jnp.transpose(UpperCAmelCase    ,  (0, 2, 3, 1)			)
     lowerCamelCase__					:   Optional[Any]      				=	self.conv_in(UpperCAmelCase			)

     # 3. down
     lowerCamelCase__					:   Any      				=	(sample,)
     for down_block in self.down_blocks:
         if isinstance(UpperCAmelCase    ,  UpperCAmelCase			):
             lowerCamelCase__				,							lowerCamelCase__					:   Optional[Any]      				=	down_block(UpperCAmelCase    ,  UpperCAmelCase    ,  UpperCAmelCase    ,  deterministic=not train			)
         else:
             lowerCamelCase__				,							lowerCamelCase__					:   Any      				=	down_block(UpperCAmelCase    ,  UpperCAmelCase    ,  deterministic=not train			)
         down_block_res_samples += res_samples

     if down_block_additional_residuals is not None:
         lowerCamelCase__					:   Union[str, Any]      				=	()

         for down_block_res_sample, down_block_additional_residual in zip(
             UpperCAmelCase    ,  UpperCAmelCase			):
             down_block_res_sample += down_block_additional_residual
             new_down_block_res_samples += (down_block_res_sample,)

         lowerCamelCase__					:   str      				=	new_down_block_res_samples

     # 4. mid
     lowerCamelCase__					:   List[Any]      				=	self.mid_block(UpperCAmelCase    ,  UpperCAmelCase    ,  UpperCAmelCase    ,  deterministic=not train			)

     if mid_block_additional_residual is not None:
         sample += mid_block_additional_residual

     # 5. up
     for up_block in self.up_blocks:
         lowerCamelCase__					:   str      				=	down_block_res_samples[-(self.layers_per_block + 1) :]
         lowerCamelCase__					:   List[str]      				=	down_block_res_samples[: -(self.layers_per_block + 1)]
         if isinstance(UpperCAmelCase    ,  UpperCAmelCase			):
             lowerCamelCase__					:   List[Any]      				=	up_block(
                 UpperCAmelCase    ,  temb=UpperCAmelCase    ,  encoder_hidden_states=UpperCAmelCase    ,  res_hidden_states_tuple=UpperCAmelCase    ,  deterministic=not train    ,  )
         else:
             lowerCamelCase__					:   int      				=	up_block(UpperCAmelCase    ,  temb=UpperCAmelCase    ,  res_hidden_states_tuple=UpperCAmelCase    ,  deterministic=not train			)

        # 6. post-process
     lowerCamelCase__					:   str      				=	self.conv_norm_out(UpperCAmelCase			)
     lowerCamelCase__					:   Optional[Any]      				=	nn.silu(UpperCAmelCase			)
     lowerCamelCase__					:   Any      				=	self.conv_out(UpperCAmelCase			)
     lowerCamelCase__					:   Optional[Any]      				=	jnp.transpose(UpperCAmelCase    ,  (0, 3, 1, 2)			)

     if not return_dict:
         return (sample,)

     return FlaxUNetaDConditionOutput(sample=UpperCAmelCase			)





