





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				(   __UpperCamelCase     ):
 UpperCAmelCase__						=						42


@flax_register_to_config
class   lowerCAmelCase				(   nn.Module,    __UpperCamelCase,    __UpperCamelCase     ):
 UpperCAmelCase__						=						32
 UpperCAmelCase__						=						4
 UpperCAmelCase__						=						4
 UpperCAmelCase__						=						(
     "CrossAttnDownBlock2D",
     "CrossAttnDownBlock2D",
     "CrossAttnDownBlock2D",
     "DownBlock2D",
 )
 UpperCAmelCase__						=						("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")
 UpperCAmelCase__						=						False
 UpperCAmelCase__						=						(3_20, 6_40, 12_80, 12_80)
 UpperCAmelCase__						=						2
 UpperCAmelCase__						=						8
 UpperCAmelCase__						=						None
 UpperCAmelCase__						=						12_80
 UpperCAmelCase__						=						0.0
 UpperCAmelCase__						=						False
 UpperCAmelCase__						=						jnp.floataa
 UpperCAmelCase__						=						True
 UpperCAmelCase__						=						0
 UpperCAmelCase__						=						False
 def 	A_      ( self    : str     ,	UpperCAmelCase    : jax.random.KeyArray      )     ->      FrozenDict:
     # init input tensors
     lowerCamelCase__       : Dict   					=						(1, self.in_channels, self.sample_size, self.sample_size)
     lowerCamelCase__       : Union[str, Any]   					=						jnp.zeros(UpperCAmelCase     ,	dtype=jnp.floataa      )
     lowerCamelCase__       : Tuple   					=						jnp.ones((1,)     ,	dtype=jnp.intaa      )
     lowerCamelCase__       : Optional[int]   					=						jnp.zeros((1, 1, self.cross_attention_dim)     ,	dtype=jnp.floataa      )

     lowerCamelCase__					,				lowerCamelCase__       : List[Any]   					=						jax.random.split(UpperCAmelCase      )
     lowerCamelCase__       : Tuple   					=						{'params': params_rng, 'dropout': dropout_rng}

     return self.init(UpperCAmelCase     ,	UpperCAmelCase     ,	UpperCAmelCase     ,	UpperCAmelCase      )["params"]
 def 	A_      ( self    : List[str]      )     ->      Union[str, Any]:
     lowerCamelCase__       : Dict   					=						self.block_out_channels
     lowerCamelCase__       : Union[str, 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.
     lowerCamelCase__       : List[Any]   					=						self.num_attention_heads or self.attention_head_dim

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

     # time
     lowerCamelCase__       : Optional[int]   					=						FlaxTimesteps(
         block_out_channels[0]     ,	flip_sin_to_cos=self.flip_sin_to_cos     ,	freq_shift=self.config.freq_shift      )
     lowerCamelCase__       : Optional[int]   					=						FlaxTimestepEmbedding(UpperCAmelCase     ,	dtype=self.dtype      )

     lowerCamelCase__       : List[str]   					=						self.only_cross_attention
     if isinstance(UpperCAmelCase     ,	UpperCAmelCase      ):
         lowerCamelCase__       : Any   					=						(only_cross_attention,) * len(self.down_block_types      )

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

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

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

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

         lowerCamelCase__       : List[str]   					=						i == len(UpperCAmelCase      ) - 1

         if up_block_type == "CrossAttnUpBlock2D":
             lowerCamelCase__       : str   					=						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:
             lowerCamelCase__       : Union[str, Any]   					=						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      )
         lowerCamelCase__       : Any   					=						output_channel
     lowerCamelCase__       : Tuple   					=						up_blocks

     # out
     lowerCamelCase__       : Union[str, Any]   					=						nn.GroupNorm(num_groups=32     ,	epsilon=1e-5      )
     lowerCamelCase__       : 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     ,	UpperCAmelCase    : int     ,	UpperCAmelCase    : Union[str, Any]     ,	UpperCAmelCase    : Union[str, Any]     ,	UpperCAmelCase    : str=None     ,	UpperCAmelCase    : Any=None     ,	UpperCAmelCase    : bool = True     ,	UpperCAmelCase    : bool = False     ,	)     ->      Union[FlaxUNetaDConditionOutput, Tuple]:
     # 1. time
     if not isinstance(UpperCAmelCase     ,	jnp.ndarray      ):
         lowerCamelCase__       : List[str]   					=						jnp.array([timesteps]     ,	dtype=jnp.intaa      )
     elif isinstance(UpperCAmelCase     ,	jnp.ndarray      ) and len(timesteps.shape      ) == 0:
         lowerCamelCase__       : List[Any]   					=						timesteps.astype(dtype=jnp.floataa      )
         lowerCamelCase__       : Optional[Any]   					=						jnp.expand_dims(UpperCAmelCase     ,	0      )

     lowerCamelCase__       : List[str]   					=						self.time_proj(UpperCAmelCase      )
     lowerCamelCase__       : Optional[Any]   					=						self.time_embedding(UpperCAmelCase      )

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

     # 3. down
     lowerCamelCase__       : Optional[Any]   					=						(sample,)
     for down_block in self.down_blocks:
         if isinstance(UpperCAmelCase     ,	UpperCAmelCase      ):
             lowerCamelCase__					,				lowerCamelCase__       : Tuple   					=						down_block(UpperCAmelCase     ,	UpperCAmelCase     ,	UpperCAmelCase     ,	deterministic=not train      )
         else:
             lowerCamelCase__					,				lowerCamelCase__       : Optional[Any]   					=						down_block(UpperCAmelCase     ,	UpperCAmelCase     ,	deterministic=not train      )
         down_block_res_samples += res_samples

     if down_block_additional_residuals is not None:
         lowerCamelCase__       : Dict   					=						()

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

         lowerCamelCase__       : int   					=						new_down_block_res_samples

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

        # 6. post-process
     lowerCamelCase__       : Dict   					=						self.conv_norm_out(UpperCAmelCase      )
     lowerCamelCase__       : str   					=						nn.silu(UpperCAmelCase      )
     lowerCamelCase__       : List[Any]   					=						self.conv_out(UpperCAmelCase      )
     lowerCamelCase__       : Dict   					=						jnp.transpose(UpperCAmelCase     ,	(0, 3, 1, 2)      )

     if not return_dict:
         return (sample,)

     return FlaxUNetaDConditionOutput(sample=UpperCAmelCase      )












from __future__ import annotations

import unittest

from transformers import is_tf_available
from transformers.testing_utils import require_sentencepiece, require_tf, require_tokenizers, slow


if is_tf_available():
    import numpy as np
    import tensorflow as tf

    from transformers import TFCamembertModel

@require_tf
@require_sentencepiece
@require_tokenizers
class   lowerCAmelCase				(   unittest.TestCase     ):
 @slow
 def 	A_      ( self    : Union[str, Any]      )     ->      List[Any]:
     lowerCamelCase__       : Tuple   					=						TFCamembertModel.from_pretrained('jplu/tf-camembert-base'      )

     lowerCamelCase__       : Any   					=						tf.convert_to_tensor(
         [[5, 121, 11, 660, 16, 730, 25543, 110, 83, 6]]     ,	dtype=tf.intaa     ,	)  # J'aime le camembert !"

     lowerCamelCase__       : List[str]   					=						model(UpperCAmelCase      )['last_hidden_state']
     lowerCamelCase__       : Tuple   					=						tf.TensorShape((1, 10, 768)      )
     self.assertEqual(output.shape     ,	UpperCAmelCase      )
     # compare the actual values for a slice.
     lowerCamelCase__       : Optional[Any]   					=						tf.convert_to_tensor(
         [[[-0.0_2_5_4, 0.0_2_3_5, 0.1_0_2_7], [0.0_6_0_6, -0.1_8_1_1, -0.0_4_1_8], [-0.1_5_6_1, -0.1_1_2_7, 0.2_6_8_7]]]     ,	dtype=tf.floataa     ,	)
     # camembert = torch.hub.load('pytorch/fairseq', 'camembert.v0')
     # camembert.eval()
     # expected_slice = roberta.model.forward(input_ids)[0][:, :3, :3].detach()

     self.assertTrue(np.allclose(output[:, :3, :3].numpy()     ,	expected_slice.numpy()     ,	atol=1e-4      )      )





