



'''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_      ):
						_lowerCAmelCase  : jnp.ndarray


@flax_register_to_config
class 						lowerCamelCase_				(       nn.Module     ,					snake_case_     ,					snake_case_      ):
						_lowerCAmelCase  : int						  =					3_2
						_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_2_0, 6_4_0, 1_2_8_0, 1_2_8_0)
						_lowerCAmelCase  : int						  =					2
						_lowerCAmelCase  : Union[int, Tuple[int]]						  =					8
						_lowerCAmelCase  : Optional[Union[int, Tuple[int]]]						  =					None
						_lowerCAmelCase  : int						  =					1_2_8_0
						_lowerCAmelCase  : float						  =					0.0
						_lowerCAmelCase  : bool						  =					False
						_lowerCAmelCase  : jnp.dtype						  =					jnp.floataa
						_lowerCAmelCase  : bool						  =					True
						_lowerCAmelCase  : int						  =					0
						_lowerCAmelCase  : bool						  =					False







						def        __lowercase     (	self :   int					,     lowerCAmelCase__ :   jax.random.KeyArray					):

													"""simple docstring"""
													# init input tensors
													SCREAMING_SNAKE_CASE       :   Optional[int]         =      (1, self.in_channels, self.sample_size, self.sample_size)
													SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      jnp.zeros(lowerCAmelCase__					,     dtype=jnp.floataa					)
													SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      jnp.ones((1,)					,     dtype=jnp.intaa					)
													SCREAMING_SNAKE_CASE       :   Any         =      jnp.zeros((1, 1, self.cross_attention_dim)					,     dtype=jnp.floataa					)

													SCREAMING_SNAKE_CASE       ,SCREAMING_SNAKE_CASE       :   Dict         =      jax.random.split(lowerCAmelCase__					)
													SCREAMING_SNAKE_CASE       :   Dict         =      {'''params''': params_rng, '''dropout''': dropout_rng}

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







						def        __lowercase     (	self :   Dict					):

													"""simple docstring"""
													SCREAMING_SNAKE_CASE       :   List[Any]         =      self.block_out_channels
													SCREAMING_SNAKE_CASE       :   Any         =      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.
													SCREAMING_SNAKE_CASE       :   Dict         =      self.num_attention_heads or self.attention_head_dim

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

													# time
													SCREAMING_SNAKE_CASE       :   Dict         =      FlaxTimesteps(
													    block_out_channels[0]					,     flip_sin_to_cos=self.flip_sin_to_cos					,     freq_shift=self.config.freq_shift					)
													SCREAMING_SNAKE_CASE       :   List[Any]         =      FlaxTimestepEmbedding(lowerCAmelCase__					,     dtype=self.dtype					)

													SCREAMING_SNAKE_CASE       :   Optional[int]         =      self.only_cross_attention
													if isinstance(lowerCAmelCase__					,     lowerCAmelCase__					):
																				SCREAMING_SNAKE_CASE       :   Tuple         =      (only_cross_attention,) * len(self.down_block_types					)

													if isinstance(lowerCAmelCase__					,     lowerCAmelCase__					):
																				SCREAMING_SNAKE_CASE       :   Optional[Any]         =      (num_attention_heads,) * len(self.down_block_types					)

													# down
													SCREAMING_SNAKE_CASE       :   Optional[int]         =      []
													SCREAMING_SNAKE_CASE       :   Any         =      block_out_channels[0]
													for i, down_block_type in enumerate(self.down_block_types					):
																				SCREAMING_SNAKE_CASE       :   Optional[int]         =      output_channel
																				SCREAMING_SNAKE_CASE       :   Tuple         =      block_out_channels[i]
																				SCREAMING_SNAKE_CASE       :   List[str]         =      i == len(lowerCAmelCase__					) - 1

																				if down_block_type == "CrossAttnDownBlock2D":
																											SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      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:
																											SCREAMING_SNAKE_CASE       :   str         =      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__					)
													SCREAMING_SNAKE_CASE       :   int         =      down_blocks

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

																				SCREAMING_SNAKE_CASE       :   int         =      i == len(lowerCAmelCase__					) - 1

																				if up_block_type == "CrossAttnUpBlock2D":
																											SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      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:
																											SCREAMING_SNAKE_CASE       :   Optional[Any]         =      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__					)
																				SCREAMING_SNAKE_CASE       :   Any         =      output_channel
													SCREAMING_SNAKE_CASE       :   Dict         =      up_blocks

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







						def __call__(	self :   Any					,     lowerCAmelCase__ :   Dict					,     lowerCAmelCase__ :   str					,     lowerCAmelCase__ :   Any					,     lowerCAmelCase__ :   Tuple=None					,     lowerCAmelCase__ :   List[str]=None					,     lowerCAmelCase__ :   bool = True					,     lowerCAmelCase__ :   bool = False					,     ):

													"""simple docstring"""
													# 1. time
													if not isinstance(lowerCAmelCase__					,     jnp.ndarray					):
																				SCREAMING_SNAKE_CASE       :   Any         =      jnp.array([timesteps]					,     dtype=jnp.intaa					)
													elif isinstance(lowerCAmelCase__					,     jnp.ndarray					) and len(timesteps.shape					) == 0:
																				SCREAMING_SNAKE_CASE       :   Dict         =      timesteps.astype(dtype=jnp.floataa					)
																				SCREAMING_SNAKE_CASE       :   int         =      jnp.expand_dims(lowerCAmelCase__					,     0					)

													SCREAMING_SNAKE_CASE       :   int         =      self.time_proj(lowerCAmelCase__					)
													SCREAMING_SNAKE_CASE       :   Dict         =      self.time_embedding(lowerCAmelCase__					)

													# 2. pre-process
													SCREAMING_SNAKE_CASE       :   List[str]         =      jnp.transpose(lowerCAmelCase__					,     (0, 2, 3, 1)					)
													SCREAMING_SNAKE_CASE       :   str         =      self.conv_in(lowerCAmelCase__					)

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

													if down_block_additional_residuals is not None:
																				SCREAMING_SNAKE_CASE       :   Optional[int]         =      ()

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

																				SCREAMING_SNAKE_CASE       :   List[str]         =      new_down_block_res_samples

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

        # 6. post-process
													SCREAMING_SNAKE_CASE       :   Optional[int]         =      self.conv_norm_out(lowerCAmelCase__					)
													SCREAMING_SNAKE_CASE       :   Any         =      nn.silu(lowerCAmelCase__					)
													SCREAMING_SNAKE_CASE       :   Dict         =      self.conv_out(lowerCAmelCase__					)
													SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      jnp.transpose(lowerCAmelCase__					,     (0, 3, 1, 2)					)

													if not return_dict:
																				return (sample,)

													return FlaxUNetaDConditionOutput(sample=lowerCAmelCase__					)









'''simple docstring'''



from typing import TYPE_CHECKING

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


lowerCAmelCase_  :							Dict             =    {
    'configuration_encodec': [
        'ENCODEC_PRETRAINED_CONFIG_ARCHIVE_MAP',
        'EncodecConfig',
    ],
    'feature_extraction_encodec': ['EncodecFeatureExtractor'],
}

try:
							if not is_torch_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							lowerCAmelCase_  :							List[Any]             =    [
							    'ENCODEC_PRETRAINED_MODEL_ARCHIVE_LIST',
							    'EncodecModel',
							    'EncodecPreTrainedModel',
							]

if TYPE_CHECKING:
							from .configuration_encodec import (
							    ENCODEC_PRETRAINED_CONFIG_ARCHIVE_MAP,
							    EncodecConfig,
							)
							from .feature_extraction_encodec import EncodecFeatureExtractor

							try:
														if not is_torch_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .modeling_encodec import (
														    ENCODEC_PRETRAINED_MODEL_ARCHIVE_LIST,
														    EncodecModel,
														    EncodecPreTrainedModel,
														)

else:
							import sys

							lowerCAmelCase_  :							List[str]             =    _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)




