
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 					_SCREAMING_SNAKE_CASE  (      __UpperCamelCase ):
      _A						:		jnp.ndarray





@flax_register_to_config
class 					_SCREAMING_SNAKE_CASE  (      nn.Module       ,					__UpperCamelCase       ,					__UpperCamelCase ):
      _A						:		int    	=     32
      _A						:		int    	=     4
      _A						:		int    	=     4
      _A						:		Tuple[str]    	=     (
          "CrossAttnDownBlock2D",
          "CrossAttnDownBlock2D",
          "CrossAttnDownBlock2D",
          "DownBlock2D",
      )
      _A						:		Tuple[str]    	=     ("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")
      _A						:		Union[bool, Tuple[bool]]    	=     False
      _A						:		Tuple[int]    	=     (3_20, 6_40, 12_80, 12_80)
      _A						:		int    	=     2
      _A						:		Union[int, Tuple[int]]    	=     8
      _A						:		Optional[Union[int, Tuple[int]]]    	=     None
      _A						:		int    	=     12_80
      _A						:		float    	=     0.0
      _A						:		bool    	=     False
      _A						:		jnp.dtype    	=     jnp.floataa
      _A						:		bool    	=     True
      _A						:		int    	=     0
      _A						:		bool    	=     False
      def A_			(			self		,	lowerCamelCase ):
          # init input tensors
          snake_case__						=				(1, self.in_channels, self.sample_size, self.sample_size)
          snake_case__						=				jnp.zeros(lowerCamelCase		,	dtype=jnp.floataa )
          snake_case__						=				jnp.ones((1,)		,	dtype=jnp.intaa )
          snake_case__						=				jnp.zeros((1, 1, self.cross_attention_dim)		,	dtype=jnp.floataa )

          snake_case__						,					snake_case__						=				jax.random.split(lowerCamelCase )
          snake_case__						=				{"params": params_rng, "dropout": dropout_rng}

          return self.init(lowerCamelCase		,	lowerCamelCase		,	lowerCamelCase		,	lowerCamelCase )["params"]
      def A_			(			self ):
          snake_case__						=				self.block_out_channels
          snake_case__						=				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.
          snake_case__						=				self.num_attention_heads or self.attention_head_dim

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

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

          snake_case__						=				self.only_cross_attention
          if isinstance(lowerCamelCase		,	lowerCamelCase ):
              snake_case__						=				(only_cross_attention,) * len(self.down_block_types )

          if isinstance(lowerCamelCase		,	lowerCamelCase ):
              snake_case__						=				(num_attention_heads,) * len(self.down_block_types )

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

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

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

              snake_case__						=				i == len(lowerCamelCase ) - 1

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

          # out
          snake_case__						=				nn.GroupNorm(num_groups=32		,	epsilon=1e-5 )
          snake_case__						=				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		,	):
          # 1. time
          if not isinstance(lowerCamelCase		,	jnp.ndarray ):
              snake_case__						=				jnp.array([timesteps]		,	dtype=jnp.intaa )
          elif isinstance(lowerCamelCase		,	jnp.ndarray ) and len(timesteps.shape ) == 0:
              snake_case__						=				timesteps.astype(dtype=jnp.floataa )
              snake_case__						=				jnp.expand_dims(lowerCamelCase		,	0 )

          snake_case__						=				self.time_proj(lowerCamelCase )
          snake_case__						=				self.time_embedding(lowerCamelCase )

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

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

          if down_block_additional_residuals is not None:
              snake_case__						=				()

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

              snake_case__						=				new_down_block_res_samples

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

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

          if not return_dict:
              return (sample,)

          return FlaxUNetaDConditionOutput(sample=lowerCamelCase )





from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
    is_vision_available,
)


__magic_name__            =  {
    '''configuration_clip''': [
        '''CLIP_PRETRAINED_CONFIG_ARCHIVE_MAP''',
        '''CLIPConfig''',
        '''CLIPOnnxConfig''',
        '''CLIPTextConfig''',
        '''CLIPVisionConfig''',
    ],
    '''processing_clip''': ['''CLIPProcessor'''],
    '''tokenization_clip''': ['''CLIPTokenizer'''],
}

try:
     if not is_tokenizers_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     __magic_name__            =  ['''CLIPTokenizerFast''']

try:
     if not is_vision_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     __magic_name__            =  ['''CLIPFeatureExtractor''']
     __magic_name__            =  ['''CLIPImageProcessor''']

try:
     if not is_torch_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     __magic_name__            =  [
         '''CLIP_PRETRAINED_MODEL_ARCHIVE_LIST''',
         '''CLIPModel''',
         '''CLIPPreTrainedModel''',
         '''CLIPTextModel''',
         '''CLIPTextModelWithProjection''',
         '''CLIPVisionModel''',
         '''CLIPVisionModelWithProjection''',
     ]

try:
     if not is_tf_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     __magic_name__            =  [
         '''TF_CLIP_PRETRAINED_MODEL_ARCHIVE_LIST''',
         '''TFCLIPModel''',
         '''TFCLIPPreTrainedModel''',
         '''TFCLIPTextModel''',
         '''TFCLIPVisionModel''',
     ]

try:
     if not is_flax_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     __magic_name__            =  [
         '''FlaxCLIPModel''',
         '''FlaxCLIPPreTrainedModel''',
         '''FlaxCLIPTextModel''',
         '''FlaxCLIPTextPreTrainedModel''',
         '''FlaxCLIPVisionModel''',
         '''FlaxCLIPVisionPreTrainedModel''',
     ]


if TYPE_CHECKING:
     from .configuration_clip import (
         CLIP_PRETRAINED_CONFIG_ARCHIVE_MAP,
         CLIPConfig,
         CLIPOnnxConfig,
         CLIPTextConfig,
         CLIPVisionConfig,
     )
     from .processing_clip import CLIPProcessor
     from .tokenization_clip import CLIPTokenizer

     try:
          if not is_tokenizers_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .tokenization_clip_fast import CLIPTokenizerFast

     try:
          if not is_vision_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .feature_extraction_clip import CLIPFeatureExtractor
          from .image_processing_clip import CLIPImageProcessor

     try:
          if not is_torch_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_clip import (
              CLIP_PRETRAINED_MODEL_ARCHIVE_LIST,
              CLIPModel,
              CLIPPreTrainedModel,
              CLIPTextModel,
              CLIPTextModelWithProjection,
              CLIPVisionModel,
              CLIPVisionModelWithProjection,
          )

     try:
          if not is_tf_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_tf_clip import (
              TF_CLIP_PRETRAINED_MODEL_ARCHIVE_LIST,
              TFCLIPModel,
              TFCLIPPreTrainedModel,
              TFCLIPTextModel,
              TFCLIPVisionModel,
          )

     try:
          if not is_flax_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_flax_clip import (
              FlaxCLIPModel,
              FlaxCLIPPreTrainedModel,
              FlaxCLIPTextModel,
              FlaxCLIPTextPreTrainedModel,
              FlaxCLIPVisionModel,
              FlaxCLIPVisionPreTrainedModel,
          )


else:
     import sys

     __magic_name__            =  _LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)



