
"""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 					__a     ( __a       ):



							'''simple docstring'''







							_lowerCamelCase					:			jnp.ndarray





@flax_register_to_config
class 					__a     ( nn.Module					,   __a					,   __a       ):



							'''simple docstring'''







							_lowerCamelCase					:			int          =   32
							_lowerCamelCase					:			int          =   4
							_lowerCamelCase					:			int          =   4
							_lowerCamelCase					:			Tuple[str]          =   (
							    "CrossAttnDownBlock2D",
							    "CrossAttnDownBlock2D",
							    "CrossAttnDownBlock2D",
							    "DownBlock2D",
							)
							_lowerCamelCase					:			Tuple[str]          =   ("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")
							_lowerCamelCase					:			Union[bool, Tuple[bool]]          =   False
							_lowerCamelCase					:			Tuple[int]          =   (3_20, 6_40, 12_80, 12_80)
							_lowerCamelCase					:			int          =   2
							_lowerCamelCase					:			Union[int, Tuple[int]]          =   8
							_lowerCamelCase					:			Optional[Union[int, Tuple[int]]]          =   None
							_lowerCamelCase					:			int          =   12_80
							_lowerCamelCase					:			float          =   0.0
							_lowerCamelCase					:			bool          =   False
							_lowerCamelCase					:			jnp.dtype          =   jnp.floataa
							_lowerCamelCase					:			bool          =   True
							_lowerCamelCase					:			int          =   0
							_lowerCamelCase					:			bool          =   False
							def   SCREAMING_SNAKE_CASE						(     self  ,    _lowerCamelCase       )			->							FrozenDict:




										'''simple docstring'''
										# init input tensors
										__lowercase       =       (1, self.in_channels, self.sample_size, self.sample_size)
										__lowercase       =       jnp.zeros(_lowerCamelCase  ,    dtype=jnp.floataa       )
										__lowercase       =       jnp.ones((1,)  ,    dtype=jnp.intaa       )
										__lowercase       =       jnp.zeros((1, 1, self.cross_attention_dim)  ,    dtype=jnp.floataa       )

										__lowercase   ,					__lowercase       =       jax.random.split(_lowerCamelCase       )
										__lowercase       =       {"params": params_rng, "dropout": dropout_rng}

										return self.init(_lowerCamelCase  ,    _lowerCamelCase  ,    _lowerCamelCase  ,    _lowerCamelCase       )["params"]
							def   SCREAMING_SNAKE_CASE						(     self       )			->							Any:




										'''simple docstring'''
										__lowercase       =       self.block_out_channels
										__lowercase       =       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.
										__lowercase       =       self.num_attention_heads or self.attention_head_dim

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

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

										__lowercase       =       self.only_cross_attention
										if isinstance(_lowerCamelCase  ,    _lowerCamelCase       ):
													__lowercase       =       (only_cross_attention,) * len(self.down_block_types       )

										if isinstance(_lowerCamelCase  ,    _lowerCamelCase       ):
													__lowercase       =       (num_attention_heads,) * len(self.down_block_types       )

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

													if down_block_type == "CrossAttnDownBlock2D":
																__lowercase       =       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:
																__lowercase       =       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       )
										__lowercase       =       down_blocks

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

													__lowercase       =       i == len(_lowerCamelCase       ) - 1

													if up_block_type == "CrossAttnUpBlock2D":
																__lowercase       =       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:
																__lowercase       =       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       )
													__lowercase       =       output_channel
										__lowercase       =       up_blocks

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





							def __call__(     self  ,    _lowerCamelCase  ,    _lowerCamelCase  ,    _lowerCamelCase  ,    _lowerCamelCase=None  ,    _lowerCamelCase=None  ,    _lowerCamelCase = True  ,    _lowerCamelCase = False  ,    )			->							Union[FlaxUNetaDConditionOutput, Tuple]:




										'''simple docstring'''
										# 1. time
										if not isinstance(_lowerCamelCase  ,    jnp.ndarray       ):
													__lowercase       =       jnp.array([timesteps]  ,    dtype=jnp.intaa       )
										elif isinstance(_lowerCamelCase  ,    jnp.ndarray       ) and len(timesteps.shape       ) == 0:
													__lowercase       =       timesteps.astype(dtype=jnp.floataa       )
													__lowercase       =       jnp.expand_dims(_lowerCamelCase  ,    0       )

										__lowercase       =       self.time_proj(_lowerCamelCase       )
										__lowercase       =       self.time_embedding(_lowerCamelCase       )

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

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

										if down_block_additional_residuals is not None:
													__lowercase       =       ()

													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,)

													__lowercase       =       new_down_block_res_samples

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

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

										if not return_dict:
													return (sample,)

										return FlaxUNetaDConditionOutput(sample=_lowerCamelCase       )







"""simple docstring"""






def      lowerCAmelCase__			(		__magic_name__ = 1_0      )						->str:
			if not isinstance(__magic_name__			,   __magic_name__      ) or n < 0:
						raise ValueError("Invalid input"      )
			__lowercase       =       1_0**n
			__lowercase       =       2_8_4_3_3 * (pow(2			,   7_8_3_0_4_5_7			,   __magic_name__      )) + 1
			return str(number % modulus      )


if __name__ == "__main__":
					from doctest import testmod

					testmod()
					print(F"{solution(10) = }")





