




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






			'''simple docstring'''



			lowerCamelCase__      =			42






@flax_register_to_config
class lowerCAmelCase__	(     nn.Module,	lowercase,	lowercase						):






			'''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 					A_    (  self  ,      lowercase  ):
							# init input tensors
							_lowerCamelCase     :			Optional[int]    =    (1, self.in_channels, self.sample_size, self.sample_size)
							_lowerCamelCase     :			Tuple    =    jnp.zeros(lowercase  ,      dtype=jnp.floataa  )
							_lowerCamelCase     :			str    =    jnp.ones((1,)  ,      dtype=jnp.intaa  )
							_lowerCamelCase     :			Optional[Any]    =    jnp.zeros((1, 1, self.cross_attention_dim)  ,      dtype=jnp.floataa  )

							_lowerCamelCase, _lowerCamelCase     :			Dict    =    jax.random.split(lowercase  )
							_lowerCamelCase     :			Union[str, Any]    =    {'params': params_rng, 'dropout': dropout_rng}

							return self.init(lowercase  ,      lowercase  ,      lowercase  ,      lowercase  )["params"]
			def 					A_    (  self  ):
							_lowerCamelCase     :			Any    =    self.block_out_channels
							_lowerCamelCase     :			int    =    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     :			str    =    self.num_attention_heads or self.attention_head_dim

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

							# time
							_lowerCamelCase     :			List[str]    =    FlaxTimesteps(
							    block_out_channels[0]  ,      flip_sin_to_cos=self.flip_sin_to_cos  ,      freq_shift=self.config.freq_shift  )
							_lowerCamelCase     :			str    =    FlaxTimestepEmbedding(lowercase  ,      dtype=self.dtype  )

							_lowerCamelCase     :			Optional[Any]    =    self.only_cross_attention
							if isinstance(lowercase  ,      lowercase  ):
											_lowerCamelCase     :			str    =    (only_cross_attention,) * len(self.down_block_types  )

							if isinstance(lowercase  ,      lowercase  ):
											_lowerCamelCase     :			Union[str, Any]    =    (num_attention_heads,) * len(self.down_block_types  )

							# down
							_lowerCamelCase     :			int    =    []
							_lowerCamelCase     :			Dict    =    block_out_channels[0]
							for i, down_block_type in enumerate(self.down_block_types  ):
											_lowerCamelCase     :			Dict    =    output_channel
											_lowerCamelCase     :			List[str]    =    block_out_channels[i]
											_lowerCamelCase     :			Union[str, Any]    =    i == len(lowercase  ) - 1

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

											down_blocks.append(lowercase  )
							_lowerCamelCase     :			int    =    down_blocks

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

											_lowerCamelCase     :			str    =    i == len(lowercase  ) - 1

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

											up_blocks.append(lowercase  )
											_lowerCamelCase     :			Dict    =    output_channel
							_lowerCamelCase     :			int    =    up_blocks

							# out
							_lowerCamelCase     :			Union[str, Any]    =    nn.GroupNorm(num_groups=32  ,      epsilon=1E-5  )
							_lowerCamelCase     :			Optional[int]    =    nn.Conv(
							    self.out_channels  ,      kernel_size=(3, 3)  ,      strides=(1, 1)  ,      padding=((1, 1), (1, 1))  ,      dtype=self.dtype  ,      )






			def __call__(  self  ,      lowercase  ,      lowercase  ,      lowercase  ,      lowercase=None  ,      lowercase=None  ,      lowercase = True  ,      lowercase = False  ,      ):
							# 1. time
							if not isinstance(lowercase  ,      jnp.ndarray  ):
											_lowerCamelCase     :			int    =    jnp.array([timesteps]  ,      dtype=jnp.intaa  )
							elif isinstance(lowercase  ,      jnp.ndarray  ) and len(timesteps.shape  ) == 0:
											_lowerCamelCase     :			Tuple    =    timesteps.astype(dtype=jnp.floataa  )
											_lowerCamelCase     :			Dict    =    jnp.expand_dims(lowercase  ,      0  )

							_lowerCamelCase     :			int    =    self.time_proj(lowercase  )
							_lowerCamelCase     :			Optional[Any]    =    self.time_embedding(lowercase  )

							# 2. pre-process
							_lowerCamelCase     :			Union[str, Any]    =    jnp.transpose(lowercase  ,      (0, 2, 3, 1)  )
							_lowerCamelCase     :			Union[str, Any]    =    self.conv_in(lowercase  )

							# 3. down
							_lowerCamelCase     :			Optional[int]    =    (sample,)
							for down_block in self.down_blocks:
											if isinstance(lowercase  ,      lowercase  ):
															_lowerCamelCase, _lowerCamelCase     :			int    =    down_block(lowercase  ,      lowercase  ,      lowercase  ,      deterministic=not train  )
											else:
															_lowerCamelCase, _lowerCamelCase     :			List[str]    =    down_block(lowercase  ,      lowercase  ,      deterministic=not train  )
											down_block_res_samples += res_samples

							if down_block_additional_residuals is not None:
											_lowerCamelCase     :			List[Any]    =    ()

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

											_lowerCamelCase     :			List[Any]    =    new_down_block_res_samples

							# 4. mid
							_lowerCamelCase     :			Tuple    =    self.mid_block(lowercase  ,      lowercase  ,      lowercase  ,      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     :			Dict    =    down_block_res_samples[-(self.layers_per_block + 1) :]
											_lowerCamelCase     :			List[Any]    =    down_block_res_samples[: -(self.layers_per_block + 1)]
											if isinstance(lowercase  ,      lowercase  ):
															_lowerCamelCase     :			str    =    up_block(
															    lowercase  ,      temb=lowercase  ,      encoder_hidden_states=lowercase  ,      res_hidden_states_tuple=lowercase  ,      deterministic=not train  ,      )
											else:
															_lowerCamelCase     :			int    =    up_block(lowercase  ,      temb=lowercase  ,      res_hidden_states_tuple=lowercase  ,      deterministic=not train  )

        # 6. post-process
							_lowerCamelCase     :			Optional[int]    =    self.conv_norm_out(lowercase  )
							_lowerCamelCase     :			Any    =    nn.silu(lowercase  )
							_lowerCamelCase     :			Optional[int]    =    self.conv_out(lowercase  )
							_lowerCamelCase     :			Any    =    jnp.transpose(lowercase  ,      (0, 3, 1, 2)  )

							if not return_dict:
											return (sample,)

							return FlaxUNetaDConditionOutput(sample=lowercase  )



"""simple docstring"""




lowercase_			     =  "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/"

def        lowercase							(		lowerCAmelCase__ :     bytes			) ->      bytes:
		# Make sure the supplied data is a bytes-like object
		if not isinstance(lowerCAmelCase__    ,					lowerCAmelCase__			):
				__a           =			f'''a bytes-like object is required, not \'{data.__class__.__name__}\''''
				raise TypeError(lowerCAmelCase__			)

		__a           =			''''''.join(bin(lowerCAmelCase__			)[2:].zfill(8			) for byte in data			)

		__a           =			len(lowerCAmelCase__			) % 6 != 0

		if padding_needed:
				# The padding that will be added later
				__a           =			b'''=''' * ((6 - len(lowerCAmelCase__			) % 6) // 2)

				# Append binary_stream with arbitrary binary digits (0's by default) to make its
				# length a multiple of 6.
				binary_stream += "0" * (6 - len(lowerCAmelCase__			) % 6)
		else:
				__a           =			b''''''

		# Encode every 6 binary digits to their corresponding Base64 character
		return (
		    "".join(
		        B64_CHARSET[int(binary_stream[index : index + 6]    ,					2			)]
		        for index in range(0    ,					len(lowerCAmelCase__			)    ,					6			)			).encode()
		    + padding
		)

def        lowercase							(		lowerCAmelCase__ :     str			) ->      bytes:
		# Make sure encoded_data is either a string or a bytes-like object
		if not isinstance(lowerCAmelCase__    ,					lowerCAmelCase__			) and not isinstance(lowerCAmelCase__    ,					lowerCAmelCase__			):
				__a           =			(
				    '''argument should be a bytes-like object or ASCII string, '''
				    f'''not \'{encoded_data.__class__.__name__}\''''
				)
				raise TypeError(lowerCAmelCase__			)

		# In case encoded_data is a bytes-like object, make sure it contains only
		# ASCII characters so we convert it to a string object
		if isinstance(lowerCAmelCase__    ,					lowerCAmelCase__			):
				try:
						__a           =			encoded_data.decode('''utf-8'''			)
				except UnicodeDecodeError:
						raise ValueError('''base64 encoded data should only contain ASCII characters'''			)

		__a           =			encoded_data.count('''='''			)

		# Check if the encoded string contains non base64 characters
		if padding:
				assert all(
				    char in B64_CHARSET for char in encoded_data[:-padding]			), "Invalid base64 character(s) found."
		else:
				assert all(
				    char in B64_CHARSET for char in encoded_data			), "Invalid base64 character(s) found."

		# Check the padding
		assert len(lowerCAmelCase__			) % 4 == 0 and padding < 3, "Incorrect padding"

		if padding:
				# Remove padding if there is one
				__a           =			encoded_data[:-padding]

				__a           =			''''''.join(
				    bin(B64_CHARSET.index(lowerCAmelCase__			)			)[2:].zfill(6			) for char in encoded_data			)[: -padding * 2]
		else:
				__a           =			''''''.join(
				    bin(B64_CHARSET.index(lowerCAmelCase__			)			)[2:].zfill(6			) for char in encoded_data			)

		__a           =			[
		    int(binary_stream[index : index + 8]    ,					2			)
		    for index in range(0    ,					len(lowerCAmelCase__			)    ,					8			)
		]

		return bytes(lowerCAmelCase__			)


if __name__ == "__main__":
				import doctest

				doctest.testmod()






