


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        UpperCamelCase__   (	__magic_name__   ):
					__SCREAMING_SNAKE_CASE       :  Optional[int]						=   42







@flax_register_to_config
class        UpperCamelCase__   (	nn.Module							,				__magic_name__							,				__magic_name__   ):
					__SCREAMING_SNAKE_CASE       :  int						=   32
					__SCREAMING_SNAKE_CASE       :  str						=   4
					__SCREAMING_SNAKE_CASE       :  str						=   4
					__SCREAMING_SNAKE_CASE       :  List[str]						=   (
					    'CrossAttnDownBlock2D',
					    'CrossAttnDownBlock2D',
					    'CrossAttnDownBlock2D',
					    'DownBlock2D',
					)
					__SCREAMING_SNAKE_CASE       :  Union[str, Any]						=   ('UpBlock2D', 'CrossAttnUpBlock2D', 'CrossAttnUpBlock2D', 'CrossAttnUpBlock2D')
					__SCREAMING_SNAKE_CASE       :  str						=   False
					__SCREAMING_SNAKE_CASE       :  str						=   (320, 640, 1280, 1280)
					__SCREAMING_SNAKE_CASE       :  Tuple						=   2
					__SCREAMING_SNAKE_CASE       :  List[Any]						=   8
					__SCREAMING_SNAKE_CASE       :  Optional[int]						=   None
					__SCREAMING_SNAKE_CASE       :  List[str]						=   1280
					__SCREAMING_SNAKE_CASE       :  str						=   0.0
					__SCREAMING_SNAKE_CASE       :  List[str]						=   False
					__SCREAMING_SNAKE_CASE       :  List[Any]						=   jnp.floataa
					__SCREAMING_SNAKE_CASE       :  str						=   True
					__SCREAMING_SNAKE_CASE       :  List[Any]						=   0
					__SCREAMING_SNAKE_CASE       :  Optional[int]						=   False

					def   UpperCAmelCase__       (     self       :   List[Any]							,					UpperCamelCase__       :   int       ):






										'''simple docstring'''






										lowercase_  							=			(1, self.in_channels, self.sample_size, self.sample_size)
										lowercase_  							=			jnp.zeros(UpperCamelCase__							,					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(UpperCamelCase__       )
										lowercase_  							=			{"""params""": params_rng, """dropout""": dropout_rng}

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

					def   UpperCAmelCase__       (     self       :   Tuple       ):






										'''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(UpperCamelCase__							,					dtype=self.dtype       )

										lowercase_  							=			self.only_cross_attention
										if isinstance(UpperCamelCase__							,					UpperCamelCase__       ):
															lowercase_  							=			(only_cross_attention,) * len(self.down_block_types       )

										if isinstance(UpperCamelCase__							,					UpperCamelCase__       ):
															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(UpperCamelCase__       ) - 1

															if down_block_type == "CrossAttnDownBlock2D":
																				lowercase_  							=			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:
																				lowercase_  							=			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__       )
										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(UpperCamelCase__       )       )
										lowercase_  							=			list(reversed(UpperCamelCase__       )       )
										lowercase_  							=			list(reversed(UpperCamelCase__       )       )
										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(UpperCamelCase__       ) - 1       )]

															lowercase_  							=			i == len(UpperCamelCase__       ) - 1

															if up_block_type == "CrossAttnUpBlock2D":
																				lowercase_  							=			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:
																				lowercase_  							=			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__       )
															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       :   str							,					UpperCamelCase__       :   Union[str, Any]							,					UpperCamelCase__       :   Dict							,					UpperCamelCase__       :   Dict							,					UpperCamelCase__       :   Any=None							,					UpperCamelCase__       :   int=None							,					UpperCamelCase__       :   Union[str, Any] = True							,					UpperCamelCase__       :   Dict = False							,					):






										'''simple docstring'''






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

										lowercase_  							=			self.time_proj(UpperCamelCase__       )
										lowercase_  							=			self.time_embedding(UpperCamelCase__       )

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

										# 3. down
										lowercase_  							=			(sample,)
										for down_block in self.down_blocks:
															if isinstance(UpperCamelCase__							,					UpperCamelCase__       ):
																				lowercase_							,					lowercase_  							=			down_block(UpperCamelCase__							,					UpperCamelCase__							,					UpperCamelCase__							,					deterministic=not train       )
															else:
																				lowercase_							,					lowercase_  							=			down_block(UpperCamelCase__							,					UpperCamelCase__							,					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(
															    UpperCamelCase__							,					UpperCamelCase__       ):
																				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(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:
															lowercase_  							=			down_block_res_samples[-(self.layers_per_block + 1) :]
															lowercase_  							=			down_block_res_samples[: -(self.layers_per_block + 1)]
															if isinstance(UpperCamelCase__							,					UpperCamelCase__       ):
																				lowercase_  							=			up_block(
																				    UpperCamelCase__							,					temb=UpperCamelCase__							,					encoder_hidden_states=UpperCamelCase__							,					res_hidden_states_tuple=UpperCamelCase__							,					deterministic=not train							,					)
															else:
																				lowercase_  							=			up_block(UpperCamelCase__							,					temb=UpperCamelCase__							,					res_hidden_states_tuple=UpperCamelCase__							,					deterministic=not train       )

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

										if not return_dict:
															return (sample,)

										return FlaxUNetaDConditionOutput(sample=UpperCamelCase__       )



def 			UpperCAmelCase_				(      UpperCAmelCase__  ):
					if not isinstance(UpperCAmelCase__						,							UpperCAmelCase__  ):
										raise TypeError("""Input value must be an 'int' type"""  )

					lowercase_  							=			0
					while number:
										position += 1
										number >>= 1

					return position


if __name__ == "__main__":
							import doctest

							doctest.testmod()


