





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

			'''simple docstring'''







			lowerCAmelCase_							  =		42






@flax_register_to_config
class       __A					(					nn.Module    ,       lowerCAmelCase    ,       lowerCAmelCase				):

			'''simple docstring'''







			lowerCAmelCase_							  =		32
			lowerCAmelCase_							  =		4
			lowerCAmelCase_							  =		4
			lowerCAmelCase_							  =		(
			    "CrossAttnDownBlock2D",
			    "CrossAttnDownBlock2D",
			    "CrossAttnDownBlock2D",
			    "DownBlock2D",
			)
			lowerCAmelCase_							  =		("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")
			lowerCAmelCase_							  =		False
			lowerCAmelCase_							  =		(320, 640, 1280, 1280)
			lowerCAmelCase_							  =		2
			lowerCAmelCase_							  =		8
			lowerCAmelCase_							  =		None
			lowerCAmelCase_							  =		1280
			lowerCAmelCase_							  =		0.0
			lowerCAmelCase_							  =		False
			lowerCAmelCase_							  =		jnp.floataa
			lowerCAmelCase_							  =		True
			lowerCAmelCase_							  =		0
			lowerCAmelCase_							  =		False





			def 	__lowerCamelCase       (					self ,						__lowerCAmelCase							):



				'''simple docstring'''







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

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

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





			def 	__lowerCamelCase       (					self							):



				'''simple docstring'''







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

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

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

				lowerCamelCase__              =	self.only_cross_attention
				if isinstance(__lowerCAmelCase ,						__lowerCAmelCase							):
					lowerCamelCase__              =	(only_cross_attention,) * len(self.down_block_types							)

				if isinstance(__lowerCAmelCase ,						__lowerCAmelCase							):
					lowerCamelCase__              =	(num_attention_heads,) * len(self.down_block_types							)

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

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

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

					lowerCamelCase__              =	i == len(__lowerCAmelCase							) - 1

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

				# out
				lowerCamelCase__              =	nn.GroupNorm(num_groups=3_2 ,						epsilon=1E-5							)
				lowerCamelCase__              =	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 ,						):



				'''simple docstring'''







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

				lowerCamelCase__              =	self.time_proj(__lowerCAmelCase							)
				lowerCamelCase__              =	self.time_embedding(__lowerCAmelCase							)

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

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

				if down_block_additional_residuals is not None:
					lowerCamelCase__              =	()

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

					lowerCamelCase__              =	new_down_block_res_samples

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

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

				if not return_dict:
					return (sample,)

				return FlaxUNetaDConditionOutput(sample=__lowerCAmelCase							)










from __future__ import annotations





def lowerCAmelCase__(__snake_case ,__snake_case = None ,__snake_case = None							)       ->   None:







	'''simple docstring'''







	if start is None:
		lowerCamelCase__              =	0

	if end is None:
		lowerCamelCase__              =	len(__snake_case							) - 1

	if start >= end:
		return

	lowerCamelCase__              =	(start + end) // 2

	slowsort(__snake_case ,__snake_case ,__snake_case							)
	slowsort(__snake_case ,mid + 1 ,__snake_case							)

	if sequence[end] < sequence[mid]:
		lowerCamelCase__       ,			lowerCamelCase__              =	sequence[mid], sequence[end]

	slowsort(__snake_case ,__snake_case ,end - 1							)


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

			testmod()



