
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					(	a__       ):
				_lowercase					=42





@flax_register_to_config
class 	_lowerCAmelCase					(	nn.Module					,      a__					,      a__       ):
				_lowercase					=32
				_lowercase					=4
				_lowercase					=4
				_lowercase					=(
				    "CrossAttnDownBlock2D",
				    "CrossAttnDownBlock2D",
				    "CrossAttnDownBlock2D",
				    "DownBlock2D",
				)
				_lowercase					=("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")
				_lowercase					=False
				_lowercase					=(3_20, 6_40, 12_80, 12_80)
				_lowercase					=2
				_lowercase					=8
				_lowercase					=None
				_lowercase					=12_80
				_lowercase					=0.0
				_lowercase					=False
				_lowercase					=jnp.floataa
				_lowercase					=True
				_lowercase					=0
				_lowercase					=False




				def      __a						(					self				,				_UpperCamelCase     )				->      FrozenDict:
							# init input tensors
							lowerCAmelCase_						=	(1, self.in_channels, self.sample_size, self.sample_size)
							lowerCAmelCase_						=	jnp.zeros(SCREAMING_SNAKE_CASE_				,				dtype=jnp.floataa     )
							lowerCAmelCase_						=	jnp.ones((1,)				,				dtype=jnp.intaa     )
							lowerCAmelCase_						=	jnp.zeros((1, 1, self.cross_attention_dim)				,				dtype=jnp.floataa     )

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

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




				def      __a						(					self     )				->      Optional[int]:
							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(SCREAMING_SNAKE_CASE_				,				dtype=self.dtype     )

							lowerCAmelCase_						=	self.only_cross_attention
							if isinstance(SCREAMING_SNAKE_CASE_				,				SCREAMING_SNAKE_CASE_     ):
										lowerCAmelCase_						=	(only_cross_attention,) * len(self.down_block_types     )

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

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

										down_blocks.append(SCREAMING_SNAKE_CASE_     )
							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(SCREAMING_SNAKE_CASE_     )     )
							lowerCAmelCase_						=	list(reversed(SCREAMING_SNAKE_CASE_     )     )
							lowerCAmelCase_						=	list(reversed(SCREAMING_SNAKE_CASE_     )     )
							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(SCREAMING_SNAKE_CASE_     ) - 1     )]

										lowerCAmelCase_						=	i == len(SCREAMING_SNAKE_CASE_     ) - 1

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

										up_blocks.append(SCREAMING_SNAKE_CASE_     )
										lowerCAmelCase_						=	output_channel
							lowerCAmelCase_						=	up_blocks

							# out
							lowerCAmelCase_						=	nn.GroupNorm(num_groups=32				,				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				,				_UpperCamelCase				,				_UpperCamelCase				,				_UpperCamelCase				,				_UpperCamelCase=None				,				_UpperCamelCase=None				,				_UpperCamelCase = True				,				_UpperCamelCase = False				,				)				->      Union[FlaxUNetaDConditionOutput, Tuple]:
							# 1. time
							if not isinstance(SCREAMING_SNAKE_CASE_				,				jnp.ndarray     ):
										lowerCAmelCase_						=	jnp.array([timesteps]				,				dtype=jnp.intaa     )
							elif isinstance(SCREAMING_SNAKE_CASE_				,				jnp.ndarray     ) and len(timesteps.shape     ) == 0:
										lowerCAmelCase_						=	timesteps.astype(dtype=jnp.floataa     )
										lowerCAmelCase_						=	jnp.expand_dims(SCREAMING_SNAKE_CASE_				,				0     )

							lowerCAmelCase_						=	self.time_proj(SCREAMING_SNAKE_CASE_     )
							lowerCAmelCase_						=	self.time_embedding(SCREAMING_SNAKE_CASE_     )

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

							# 3. down
							lowerCAmelCase_						=	(sample,)
							for down_block in self.down_blocks:
										if isinstance(SCREAMING_SNAKE_CASE_				,				SCREAMING_SNAKE_CASE_     ):
													lowerCAmelCase_						=	down_block(SCREAMING_SNAKE_CASE_				,				SCREAMING_SNAKE_CASE_				,				SCREAMING_SNAKE_CASE_				,				deterministic=not train     )
										else:
													lowerCAmelCase_						=	down_block(SCREAMING_SNAKE_CASE_				,				SCREAMING_SNAKE_CASE_				,				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(
										    SCREAMING_SNAKE_CASE_				,				SCREAMING_SNAKE_CASE_     ):
													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(SCREAMING_SNAKE_CASE_				,				SCREAMING_SNAKE_CASE_				,				SCREAMING_SNAKE_CASE_				,				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(SCREAMING_SNAKE_CASE_				,				SCREAMING_SNAKE_CASE_     ):
													lowerCAmelCase_						=	up_block(
													    SCREAMING_SNAKE_CASE_				,				temb=SCREAMING_SNAKE_CASE_				,				encoder_hidden_states=SCREAMING_SNAKE_CASE_				,				res_hidden_states_tuple=SCREAMING_SNAKE_CASE_				,				deterministic=not train				,				)
										else:
													lowerCAmelCase_						=	up_block(SCREAMING_SNAKE_CASE_				,				temb=SCREAMING_SNAKE_CASE_				,				res_hidden_states_tuple=SCREAMING_SNAKE_CASE_				,				deterministic=not train     )

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

							if not return_dict:
										return (sample,)

							return FlaxUNetaDConditionOutput(sample=SCREAMING_SNAKE_CASE_     )








from collections import OrderedDict
from typing import TYPE_CHECKING, Any, List, Mapping, Optional

from packaging import version


if TYPE_CHECKING:
    from ... import PreTrainedTokenizer, TensorType

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfigWithPast, PatchingSpec
from ...utils import is_torch_available, logging


__UpperCAmelCase          =       logging.get_logger(__name__)

__UpperCAmelCase          =       {
    '''bigscience/bloom''': '''https://huggingface.co/bigscience/bloom/resolve/main/config.json''',
    '''bigscience/bloom-560m''': '''https://huggingface.co/bigscience/bloom-560m/blob/main/config.json''',
    '''bigscience/bloom-1b1''': '''https://huggingface.co/bigscience/bloom-1b1/blob/main/config.json''',
    '''bigscience/bloom-1b7''': '''https://huggingface.co/bigscience/bloom-1b7/blob/main/config.json''',
    '''bigscience/bloom-3b''': '''https://huggingface.co/bigscience/bloom-3b/blob/main/config.json''',
    '''bigscience/bloom-7b1''': '''https://huggingface.co/bigscience/bloom-7b1/blob/main/config.json''',
}



class 				lowerCAmelCase_		(  a__							):
       UpperCAmelCase__				:					Tuple        =  "bloom"
       UpperCAmelCase__				:					List[Any]        =  ["past_key_values"]
       UpperCAmelCase__				:					str        =  {
           "num_hidden_layers": "n_layer",
           "num_attention_heads": "n_head",
       }
       def __init__(			self,			SCREAMING_SNAKE_CASE_=25_0880,			SCREAMING_SNAKE_CASE_=64,			SCREAMING_SNAKE_CASE_=2,			SCREAMING_SNAKE_CASE_=8,			SCREAMING_SNAKE_CASE_=1e-5,			SCREAMING_SNAKE_CASE_=0.02,			SCREAMING_SNAKE_CASE_=True,			SCREAMING_SNAKE_CASE_=1,			SCREAMING_SNAKE_CASE_=2,			SCREAMING_SNAKE_CASE_=False,			SCREAMING_SNAKE_CASE_=0.0,			SCREAMING_SNAKE_CASE_=0.0,			SCREAMING_SNAKE_CASE_=1,			SCREAMING_SNAKE_CASE_=False,			**SCREAMING_SNAKE_CASE_,			)     ->		Tuple:
          UpperCamelCase    :  str							=	vocab_size
          # Backward compatibility with n_embed kwarg
          UpperCamelCase    :  Optional[Any]							=	kwargs.pop('n_embed',			SCREAMING_SNAKE_CASE_       )
          UpperCamelCase    :  Union[str, Any]							=	hidden_size if n_embed is None else n_embed
          UpperCamelCase    :  Tuple							=	n_layer
          UpperCamelCase    :  Dict							=	n_head
          UpperCamelCase    :  List[Any]							=	layer_norm_epsilon
          UpperCamelCase    :  Optional[int]							=	initializer_range
          UpperCamelCase    :  int							=	use_cache
          UpperCamelCase    :  int							=	pretraining_tp
          UpperCamelCase    :  Optional[int]							=	apply_residual_connection_post_layernorm
          UpperCamelCase    :  str							=	hidden_dropout
          UpperCamelCase    :  str							=	attention_dropout

          UpperCamelCase    :  List[Any]							=	bos_token_id
          UpperCamelCase    :  Tuple							=	eos_token_id
          UpperCamelCase    :  Union[str, Any]							=	slow_but_exact

          super().__init__(bos_token_id=SCREAMING_SNAKE_CASE_,			eos_token_id=SCREAMING_SNAKE_CASE_,			**SCREAMING_SNAKE_CASE_       )






class 				lowerCAmelCase_		(  a__							):
       UpperCAmelCase__				:					Any        =  version.parse("1.12"							)
       def __init__(			self,			SCREAMING_SNAKE_CASE_,			SCREAMING_SNAKE_CASE_ = "default",			SCREAMING_SNAKE_CASE_ = None,			SCREAMING_SNAKE_CASE_ = False,			)     ->		Any:
          super().__init__(SCREAMING_SNAKE_CASE_,			task=SCREAMING_SNAKE_CASE_,			patching_specs=SCREAMING_SNAKE_CASE_,			use_past=SCREAMING_SNAKE_CASE_       )
          if not getattr(self._config,			'pad_token_id',			SCREAMING_SNAKE_CASE_       ):
             # TODO: how to do that better?
             UpperCamelCase    :  Tuple							=	0
       @property
       def 					snake_case_       (			self       )     ->		Mapping[str, Mapping[int, str]]:
          UpperCamelCase    :  str							=	OrderedDict({'input_ids': {0: 'batch', 1: 'sequence'}}       )
          if self.use_past:
             # BLOOM stores values on dynamic axis 2. For more details see: https://github.com/huggingface/transformers/pull/18344
             self.fill_with_past_key_values_(SCREAMING_SNAKE_CASE_,			direction='inputs',			inverted_values_shape=SCREAMING_SNAKE_CASE_       )
             UpperCamelCase    :  Union[str, Any]							=	{0: 'batch', 1: 'past_sequence + sequence'}
          else:
             UpperCamelCase    :  Optional[int]							=	{0: 'batch', 1: 'sequence'}

          return common_inputs
       @property
       def 					snake_case_       (			self       )     ->		int:
          return self._config.n_layer
       @property
       def 					snake_case_       (			self       )     ->		int:
          return self._config.n_head
       @property
       def 					snake_case_       (			self       )     ->		float:
          return 1e-3
       def 					snake_case_       (			self,			SCREAMING_SNAKE_CASE_,			SCREAMING_SNAKE_CASE_ = -1,			SCREAMING_SNAKE_CASE_ = -1,			SCREAMING_SNAKE_CASE_ = False,			SCREAMING_SNAKE_CASE_ = None,			)     ->		Mapping[str, Any]:
          UpperCamelCase    :  Dict							=	super(SCREAMING_SNAKE_CASE_,			self       ).generate_dummy_inputs(
              SCREAMING_SNAKE_CASE_,			batch_size=SCREAMING_SNAKE_CASE_,			seq_length=SCREAMING_SNAKE_CASE_,			is_pair=SCREAMING_SNAKE_CASE_,			framework=SCREAMING_SNAKE_CASE_       )

          # We need to order the input in the way they appears in the forward()
          UpperCamelCase    :  Any							=	OrderedDict({'input_ids': common_inputs['input_ids']}       )

          # Need to add the past_keys
          if self.use_past:
             if not is_torch_available():
                raise ValueError('Cannot generate dummy past_keys inputs without PyTorch installed.'       )
             else:
                import torch

                UpperCamelCase	,     UpperCamelCase    :  int							=	common_inputs['input_ids'].shape
                # Not using the same length for past_key_values
                UpperCamelCase    :  Any							=	seqlen + 2
                UpperCamelCase    :  Optional[int]							=	self._config.hidden_size // self.num_attention_heads
                UpperCamelCase    :  Any							=	(
                    batch * self.num_attention_heads,
                    head_dim,
                    past_key_values_length,
                )
                UpperCamelCase    :  Optional[Any]							=	(
                    batch * self.num_attention_heads,
                    past_key_values_length,
                    head_dim,
                )
                UpperCamelCase    :  List[str]							=	[
                    (torch.zeros(SCREAMING_SNAKE_CASE_       ), torch.zeros(SCREAMING_SNAKE_CASE_       )) for _ in range(self.num_layers       )
                ]

          UpperCamelCase    :  str							=	common_inputs['attention_mask']
          if self.use_past:
             UpperCamelCase    :  int							=	ordered_inputs['attention_mask'].dtype
             UpperCamelCase    :  List[Any]							=	torch.cat(
                 [ordered_inputs['attention_mask'], torch.ones(SCREAMING_SNAKE_CASE_,			SCREAMING_SNAKE_CASE_,			dtype=SCREAMING_SNAKE_CASE_       )],			dim=1       )

          return ordered_inputs






       @property
       def 					snake_case_       (			self       )     ->		int:
          return 13



