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







							'''simple docstring'''

							lowerCAmelCase_ 						=						42


@flax_register_to_config
class       _lowerCAmelCase   (   nn.Module					,							__snake_case					,							__snake_case				):







							'''simple docstring'''

							lowerCAmelCase_ 						=						32
							lowerCAmelCase_ 						=						4
							lowerCAmelCase_ 						=						4
							lowerCAmelCase_ 						=						(
							    "CrossAttnDownBlock2D",
							    "CrossAttnDownBlock2D",
							    "CrossAttnDownBlock2D",
							    "DownBlock2D",
							)
							lowerCAmelCase_ 						=						("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")
							lowerCAmelCase_ 						=						False
							lowerCAmelCase_ 						=						(3_20, 6_40, 12_80, 12_80)
							lowerCAmelCase_ 						=						2
							lowerCAmelCase_ 						=						8
							lowerCAmelCase_ 						=						None
							lowerCAmelCase_ 						=						12_80
							lowerCAmelCase_ 						=						0.0
							lowerCAmelCase_ 						=						False
							lowerCAmelCase_ 						=						jnp.floataa
							lowerCAmelCase_ 						=						True
							lowerCAmelCase_ 						=						0
							lowerCAmelCase_ 						=						False

							def 	lowercase							(self  ,  UpperCAmelCase  )					->       FrozenDict:
														# init input tensors
														_snake_case       =  (1, self.in_channels, self.sample_size, self.sample_size)
														_snake_case       =  jnp.zeros(UpperCAmelCase  ,  dtype=jnp.floataa  )
														_snake_case       =  jnp.ones((1,)  ,  dtype=jnp.intaa  )
														_snake_case       =  jnp.zeros((1, 1, self.cross_attention_dim)  ,  dtype=jnp.floataa  )

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

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

							def 	lowercase							(self  )					->       Any:
														_snake_case       =  self.block_out_channels
														_snake_case       =  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.
														_snake_case       =  self.num_attention_heads or self.attention_head_dim

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

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

														_snake_case       =  self.only_cross_attention
														if isinstance(UpperCAmelCase  ,  UpperCAmelCase  ):
																					_snake_case       =  (only_cross_attention,) * len(self.down_block_types  )

														if isinstance(UpperCAmelCase  ,  UpperCAmelCase  ):
																					_snake_case       =  (num_attention_heads,) * len(self.down_block_types  )

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

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

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

																					_snake_case       =  i == len(UpperCAmelCase  ) - 1

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

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





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

														_snake_case       =  self.time_proj(UpperCAmelCase  )
														_snake_case       =  self.time_embedding(UpperCAmelCase  )

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

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

														if down_block_additional_residuals is not None:
																					_snake_case       =  ()

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

																					_snake_case       =  new_down_block_res_samples

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

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

														if not return_dict:
																					return (sample,)

														return FlaxUNetaDConditionOutput(sample=UpperCAmelCase  )

'''simple docstring'''
from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_tf_available,
    is_torch_available,
    is_vision_available,
)


__lowerCAmelCase												=	{
    'configuration_efficientformer': [
        'EFFICIENTFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP',
        'EfficientFormerConfig',
    ]
}

try:
				if not is_vision_available():
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				pass
else:
				__lowerCAmelCase												=	['EfficientFormerImageProcessor']

try:
				if not is_torch_available():
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				pass
else:
				__lowerCAmelCase												=	[
				    'EFFICIENTFORMER_PRETRAINED_MODEL_ARCHIVE_LIST',
				    'EfficientFormerForImageClassification',
				    'EfficientFormerForImageClassificationWithTeacher',
				    'EfficientFormerModel',
				    'EfficientFormerPreTrainedModel',
				]

try:
				if not is_tf_available():
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				pass
else:
				__lowerCAmelCase												=	[
				    'TF_EFFICIENTFORMER_PRETRAINED_MODEL_ARCHIVE_LIST',
				    'TFEfficientFormerForImageClassification',
				    'TFEfficientFormerForImageClassificationWithTeacher',
				    'TFEfficientFormerModel',
				    'TFEfficientFormerPreTrainedModel',
				]

if TYPE_CHECKING:
				from .configuration_efficientformer import EFFICIENTFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP, EfficientFormerConfig

				try:
								if not is_vision_available():
												raise OptionalDependencyNotAvailable()
				except OptionalDependencyNotAvailable:
								pass
				else:
								from .image_processing_efficientformer import EfficientFormerImageProcessor

				try:
								if not is_torch_available():
												raise OptionalDependencyNotAvailable()
				except OptionalDependencyNotAvailable:
								pass
				else:
								from .modeling_efficientformer import (
								    EFFICIENTFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
								    EfficientFormerForImageClassification,
								    EfficientFormerForImageClassificationWithTeacher,
								    EfficientFormerModel,
								    EfficientFormerPreTrainedModel,
								)
				try:
								if not is_tf_available():
												raise OptionalDependencyNotAvailable()
				except OptionalDependencyNotAvailable:
								pass
				else:
								from .modeling_tf_efficientformer import (
								    TF_EFFICIENTFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
								    TFEfficientFormerForImageClassification,
								    TFEfficientFormerForImageClassificationWithTeacher,
								    TFEfficientFormerModel,
								    TFEfficientFormerPreTrainedModel,
								)

else:
				import sys

				__lowerCAmelCase												=	_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)