





'''simple docstring'''

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	(       __magic_name__					):






       """simple docstring"""
       _snake_case      :	jnp.ndarray

@flax_register_to_config
class 							__lowerCAmelCase	(       nn.Module		,		__magic_name__		,		__magic_name__					):






       """simple docstring"""
       _snake_case      :	int									=       3_2
       _snake_case      :	int									=       4
       _snake_case      :	int									=       4
       _snake_case      :	Tuple[str]									=       (
           "CrossAttnDownBlock2D",
           "CrossAttnDownBlock2D",
           "CrossAttnDownBlock2D",
           "DownBlock2D",
       )
       _snake_case      :	Tuple[str]									=       ("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")
       _snake_case      :	Union[bool, Tuple[bool]]									=       False
       _snake_case      :	Tuple[int]									=       (3_2_0, 6_4_0, 1_2_8_0, 1_2_8_0)
       _snake_case      :	int									=       2
       _snake_case      :	Union[int, Tuple[int]]									=       8
       _snake_case      :	Optional[Union[int, Tuple[int]]]									=       None
       _snake_case      :	int									=       1_2_8_0
       _snake_case      :	float									=       0.0
       _snake_case      :	bool									=       False
       _snake_case      :	jnp.dtype									=       jnp.floataa
       _snake_case      :	bool									=       True
       _snake_case      :	int									=       0
       _snake_case      :	bool									=       False





       def snake_case__		(				self  :						str			,					lowerCAmelCase__  :						jax.random.KeyArray	)       ->	FrozenDict:

            '''simple docstring'''




            _UpperCamelCase							       =			(1, self.in_channels, self.sample_size, self.sample_size)
            _UpperCamelCase							       =			jnp.zeros(lowerCAmelCase__			,					dtype=jnp.floataa	)
            _UpperCamelCase							       =			jnp.ones((1,)			,					dtype=jnp.intaa	)
            _UpperCamelCase							       =			jnp.zeros((1, 1, self.cross_attention_dim)			,					dtype=jnp.floataa	)

            _UpperCamelCase       ,							_UpperCamelCase							       =			jax.random.split(lowerCAmelCase__	)
            _UpperCamelCase							       =			{'''params''': params_rng, '''dropout''': dropout_rng}

            return self.init(lowerCAmelCase__			,					lowerCAmelCase__			,					lowerCAmelCase__			,					lowerCAmelCase__	)["params"]





       def snake_case__		(				self  :						str	)       ->	Tuple:

            '''simple docstring'''




            _UpperCamelCase							       =			self.block_out_channels
            _UpperCamelCase							       =			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.
            _UpperCamelCase							       =			self.num_attention_heads or self.attention_head_dim

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

            # time
            _UpperCamelCase							       =			FlaxTimesteps(
                block_out_channels[0]			,					flip_sin_to_cos=self.flip_sin_to_cos			,					freq_shift=self.config.freq_shift	)
            _UpperCamelCase							       =			FlaxTimestepEmbedding(lowerCAmelCase__			,					dtype=self.dtype	)

            _UpperCamelCase							       =			self.only_cross_attention
            if isinstance(lowerCAmelCase__			,					lowerCAmelCase__	):
                 _UpperCamelCase							       =			(only_cross_attention,) * len(self.down_block_types	)

            if isinstance(lowerCAmelCase__			,					lowerCAmelCase__	):
                 _UpperCamelCase							       =			(num_attention_heads,) * len(self.down_block_types	)

            # down
            _UpperCamelCase							       =			[]
            _UpperCamelCase							       =			block_out_channels[0]
            for i, down_block_type in enumerate(self.down_block_types	):
                 _UpperCamelCase							       =			output_channel
                 _UpperCamelCase							       =			block_out_channels[i]
                 _UpperCamelCase							       =			i == len(lowerCAmelCase__	) - 1

                 if down_block_type == "CrossAttnDownBlock2D":
                      _UpperCamelCase							       =			FlaxCrossAttnDownBlockaD(
                          in_channels=lowerCAmelCase__			,					out_channels=lowerCAmelCase__			,					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:
                      _UpperCamelCase							       =			FlaxDownBlockaD(
                          in_channels=lowerCAmelCase__			,					out_channels=lowerCAmelCase__			,					dropout=self.dropout			,					num_layers=self.layers_per_block			,					add_downsample=not is_final_block			,					dtype=self.dtype			,					)

                 down_blocks.append(lowerCAmelCase__	)
            _UpperCamelCase							       =			down_blocks

            # mid
            _UpperCamelCase							       =			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
            _UpperCamelCase							       =			[]
            _UpperCamelCase							       =			list(reversed(lowerCAmelCase__	)	)
            _UpperCamelCase							       =			list(reversed(lowerCAmelCase__	)	)
            _UpperCamelCase							       =			list(reversed(lowerCAmelCase__	)	)
            _UpperCamelCase							       =			reversed_block_out_channels[0]
            for i, up_block_type in enumerate(self.up_block_types	):
                 _UpperCamelCase							       =			output_channel
                 _UpperCamelCase							       =			reversed_block_out_channels[i]
                 _UpperCamelCase							       =			reversed_block_out_channels[min(i + 1			,					len(lowerCAmelCase__	) - 1	)]

                 _UpperCamelCase							       =			i == len(lowerCAmelCase__	) - 1

                 if up_block_type == "CrossAttnUpBlock2D":
                      _UpperCamelCase							       =			FlaxCrossAttnUpBlockaD(
                          in_channels=lowerCAmelCase__			,					out_channels=lowerCAmelCase__			,					prev_output_channel=lowerCAmelCase__			,					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:
                      _UpperCamelCase							       =			FlaxUpBlockaD(
                          in_channels=lowerCAmelCase__			,					out_channels=lowerCAmelCase__			,					prev_output_channel=lowerCAmelCase__			,					num_layers=self.layers_per_block + 1			,					add_upsample=not is_final_block			,					dropout=self.dropout			,					dtype=self.dtype			,					)

                 up_blocks.append(lowerCAmelCase__	)
                 _UpperCamelCase							       =			output_channel
            _UpperCamelCase							       =			up_blocks

            # out
            _UpperCamelCase							       =			nn.GroupNorm(num_groups=32			,					epsilon=1e-5	)
            _UpperCamelCase							       =			nn.Conv(
                self.out_channels			,					kernel_size=(3, 3)			,					strides=(1, 1)			,					padding=((1, 1), (1, 1))			,					dtype=self.dtype			,					)







       def __call__(				self  :						Tuple			,					lowerCAmelCase__  :						Dict			,					lowerCAmelCase__  :						List[Any]			,					lowerCAmelCase__  :						Dict			,					lowerCAmelCase__  :						Tuple=None			,					lowerCAmelCase__  :						Union[str, Any]=None			,					lowerCAmelCase__  :						bool = True			,					lowerCAmelCase__  :						bool = False			,					)       ->	Union[FlaxUNetaDConditionOutput, Tuple]:

            '''simple docstring'''




            if not isinstance(lowerCAmelCase__			,					jnp.ndarray	):
                 _UpperCamelCase							       =			jnp.array([timesteps]			,					dtype=jnp.intaa	)
            elif isinstance(lowerCAmelCase__			,					jnp.ndarray	) and len(timesteps.shape	) == 0:
                 _UpperCamelCase							       =			timesteps.astype(dtype=jnp.floataa	)
                 _UpperCamelCase							       =			jnp.expand_dims(lowerCAmelCase__			,					0	)

            _UpperCamelCase							       =			self.time_proj(lowerCAmelCase__	)
            _UpperCamelCase							       =			self.time_embedding(lowerCAmelCase__	)

            # 2. pre-process
            _UpperCamelCase							       =			jnp.transpose(lowerCAmelCase__			,					(0, 2, 3, 1)	)
            _UpperCamelCase							       =			self.conv_in(lowerCAmelCase__	)

            # 3. down
            _UpperCamelCase							       =			(sample,)
            for down_block in self.down_blocks:
                 if isinstance(lowerCAmelCase__			,					lowerCAmelCase__	):
                      _UpperCamelCase       ,							_UpperCamelCase							       =			down_block(lowerCAmelCase__			,					lowerCAmelCase__			,					lowerCAmelCase__			,					deterministic=not train	)
                 else:
                      _UpperCamelCase       ,							_UpperCamelCase							       =			down_block(lowerCAmelCase__			,					lowerCAmelCase__			,					deterministic=not train	)
                 down_block_res_samples += res_samples

            if down_block_additional_residuals is not None:
                 _UpperCamelCase							       =			()

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

                 _UpperCamelCase							       =			new_down_block_res_samples

            # 4. mid
            _UpperCamelCase							       =			self.mid_block(lowerCAmelCase__			,					lowerCAmelCase__			,					lowerCAmelCase__			,					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:
                 _UpperCamelCase							       =			down_block_res_samples[-(self.layers_per_block + 1) :]
                 _UpperCamelCase							       =			down_block_res_samples[: -(self.layers_per_block + 1)]
                 if isinstance(lowerCAmelCase__			,					lowerCAmelCase__	):
                      _UpperCamelCase							       =			up_block(
                          lowerCAmelCase__			,					temb=lowerCAmelCase__			,					encoder_hidden_states=lowerCAmelCase__			,					res_hidden_states_tuple=lowerCAmelCase__			,					deterministic=not train			,					)
                 else:
                      _UpperCamelCase							       =			up_block(lowerCAmelCase__			,					temb=lowerCAmelCase__			,					res_hidden_states_tuple=lowerCAmelCase__			,					deterministic=not train	)

        # 6. post-process
            _UpperCamelCase							       =			self.conv_norm_out(lowerCAmelCase__	)
            _UpperCamelCase							       =			nn.silu(lowerCAmelCase__	)
            _UpperCamelCase							       =			self.conv_out(lowerCAmelCase__	)
            _UpperCamelCase							       =			jnp.transpose(lowerCAmelCase__			,					(0, 3, 1, 2)	)

            if not return_dict:
                 return (sample,)

            return FlaxUNetaDConditionOutput(sample=lowerCAmelCase__	)







"""simple docstring"""



from diffusers.utils.testing_utils import require_onnxruntime







@require_onnxruntime
class    __a						:
				pass