

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






@flax.struct.dataclass
class 	UpperCAmelCase__		(	__UpperCamelCase	):






  """simple docstring"""


  a									= 42
  a									= 42






class 	UpperCAmelCase__		(	nn.Module	):






  """simple docstring"""


  a									= 42
  a									= (16, 32, 96, 2_56)
  a									= jnp.floataa






  def 				lowercase_		(							self						:		str )		->							Optional[int]:
       SCREAMING_SNAKE_CASE__   =   nn.Conv(
           self.block_out_channels[0]     ,       kernel_size=(3, 3)     ,       padding=((1, 1), (1, 1))     ,       dtype=self.dtype     ,       )

       SCREAMING_SNAKE_CASE__   =   []
       for i in range(len(self.block_out_channels ) - 1 ):
            SCREAMING_SNAKE_CASE__   =   self.block_out_channels[i]
            SCREAMING_SNAKE_CASE__   =   self.block_out_channels[i + 1]
            SCREAMING_SNAKE_CASE__   =   nn.Conv(
                __lowerCamelCase     ,       kernel_size=(3, 3)     ,       padding=((1, 1), (1, 1))     ,       dtype=self.dtype     ,       )
            blocks.append(__lowerCamelCase )
            SCREAMING_SNAKE_CASE__   =   nn.Conv(
                __lowerCamelCase     ,       kernel_size=(3, 3)     ,       strides=(2, 2)     ,       padding=((1, 1), (1, 1))     ,       dtype=self.dtype     ,       )
            blocks.append(__lowerCamelCase )
       SCREAMING_SNAKE_CASE__   =   blocks

       SCREAMING_SNAKE_CASE__   =   nn.Conv(
           self.conditioning_embedding_channels     ,       kernel_size=(3, 3)     ,       padding=((1, 1), (1, 1))     ,       kernel_init=nn.initializers.zeros_init()     ,       bias_init=nn.initializers.zeros_init()     ,       dtype=self.dtype     ,       )







  def __call__(							self						:		Optional[int]     ,       __lowerCamelCase						:		List[str] )		->							int:
       SCREAMING_SNAKE_CASE__   =   self.conv_in(__lowerCamelCase )
       SCREAMING_SNAKE_CASE__   =   nn.silu(__lowerCamelCase )

       for block in self.blocks:
            SCREAMING_SNAKE_CASE__   =   block(__lowerCamelCase )
            SCREAMING_SNAKE_CASE__   =   nn.silu(__lowerCamelCase )

       SCREAMING_SNAKE_CASE__   =   self.conv_out(__lowerCamelCase )

       return embedding






@flax_register_to_config
class 	UpperCAmelCase__		(	nn.Module						,					__UpperCamelCase						,					__UpperCamelCase	):






  """simple docstring"""


  a									= 32
  a									= 4
  a									= (
      "CrossAttnDownBlock2D",
      "CrossAttnDownBlock2D",
      "CrossAttnDownBlock2D",
      "DownBlock2D",
  )
  a									= False
  a									= (3_20, 6_40, 12_80, 12_80)
  a									= 2
  a									= 8
  a									= None
  a									= 12_80
  a									= 0.0
  a									= False
  a									= jnp.floataa
  a									= True
  a									= 0
  a									= "rgb"
  a									= (16, 32, 96, 2_56)






  def 				lowercase_		(							self						:		int     ,       __lowerCamelCase						:		jax.random.KeyArray )		->							Tuple:
       # init input tensors
       SCREAMING_SNAKE_CASE__   =   (1, self.in_channels, self.sample_size, self.sample_size)
       SCREAMING_SNAKE_CASE__   =   jnp.zeros(__lowerCamelCase     ,       dtype=jnp.floataa )
       SCREAMING_SNAKE_CASE__   =   jnp.ones((1,)     ,       dtype=jnp.intaa )
       SCREAMING_SNAKE_CASE__   =   jnp.zeros((1, 1, self.cross_attention_dim)     ,       dtype=jnp.floataa )
       SCREAMING_SNAKE_CASE__   =   (1, 3, self.sample_size * 8, self.sample_size * 8)
       SCREAMING_SNAKE_CASE__   =   jnp.zeros(__lowerCamelCase     ,       dtype=jnp.floataa )

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

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






  def 				lowercase_		(							self						:		Tuple )		->							str:
       SCREAMING_SNAKE_CASE__   =   self.block_out_channels
       SCREAMING_SNAKE_CASE__   =   block_out_channels[0] * 4

       # 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__   =   self.num_attention_heads or self.attention_head_dim

       # input
       SCREAMING_SNAKE_CASE__   =   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__   =   FlaxTimesteps(
           block_out_channels[0]     ,       flip_sin_to_cos=self.flip_sin_to_cos     ,       freq_shift=self.config.freq_shift )
       SCREAMING_SNAKE_CASE__   =   FlaxTimestepEmbedding(__lowerCamelCase     ,       dtype=self.dtype )

       SCREAMING_SNAKE_CASE__   =   FlaxControlNetConditioningEmbedding(
           conditioning_embedding_channels=block_out_channels[0]     ,       block_out_channels=self.conditioning_embedding_out_channels     ,       )

       SCREAMING_SNAKE_CASE__   =   self.only_cross_attention
       if isinstance(__lowerCamelCase     ,       __lowerCamelCase ):
            SCREAMING_SNAKE_CASE__   =   (only_cross_attention,) * len(self.down_block_types )

       if isinstance(__lowerCamelCase     ,       __lowerCamelCase ):
            SCREAMING_SNAKE_CASE__   =   (num_attention_heads,) * len(self.down_block_types )

       # down
       SCREAMING_SNAKE_CASE__   =   []
       SCREAMING_SNAKE_CASE__   =   []

       SCREAMING_SNAKE_CASE__   =   block_out_channels[0]

       SCREAMING_SNAKE_CASE__   =   nn.Conv(
           __lowerCamelCase     ,       kernel_size=(1, 1)     ,       padding='''VALID'''     ,       kernel_init=nn.initializers.zeros_init()     ,       bias_init=nn.initializers.zeros_init()     ,       dtype=self.dtype     ,       )
       controlnet_down_blocks.append(__lowerCamelCase )

       for i, down_block_type in enumerate(self.down_block_types ):
            SCREAMING_SNAKE_CASE__   =   output_channel
            SCREAMING_SNAKE_CASE__   =   block_out_channels[i]
            SCREAMING_SNAKE_CASE__   =   i == len(__lowerCamelCase ) - 1

            if down_block_type == "CrossAttnDownBlock2D":
                 SCREAMING_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]     ,       dtype=self.dtype     ,       )
            else:
                 SCREAMING_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 )

            for _ in range(self.layers_per_block ):
                 SCREAMING_SNAKE_CASE__   =   nn.Conv(
                     __lowerCamelCase     ,       kernel_size=(1, 1)     ,       padding='''VALID'''     ,       kernel_init=nn.initializers.zeros_init()     ,       bias_init=nn.initializers.zeros_init()     ,       dtype=self.dtype     ,       )
                 controlnet_down_blocks.append(__lowerCamelCase )

            if not is_final_block:
                 SCREAMING_SNAKE_CASE__   =   nn.Conv(
                     __lowerCamelCase     ,       kernel_size=(1, 1)     ,       padding='''VALID'''     ,       kernel_init=nn.initializers.zeros_init()     ,       bias_init=nn.initializers.zeros_init()     ,       dtype=self.dtype     ,       )
                 controlnet_down_blocks.append(__lowerCamelCase )

       SCREAMING_SNAKE_CASE__   =   down_blocks
       SCREAMING_SNAKE_CASE__   =   controlnet_down_blocks

       # mid
       SCREAMING_SNAKE_CASE__   =   block_out_channels[-1]
       SCREAMING_SNAKE_CASE__   =   FlaxUNetMidBlockaDCrossAttn(
           in_channels=__lowerCamelCase     ,       dropout=self.dropout     ,       num_attention_heads=num_attention_heads[-1]     ,       use_linear_projection=self.use_linear_projection     ,       dtype=self.dtype     ,       )

       SCREAMING_SNAKE_CASE__   =   nn.Conv(
           __lowerCamelCase     ,       kernel_size=(1, 1)     ,       padding='''VALID'''     ,       kernel_init=nn.initializers.zeros_init()     ,       bias_init=nn.initializers.zeros_init()     ,       dtype=self.dtype     ,       )







  def __call__(							self						:		List[str]     ,       __lowerCamelCase						:		Union[str, Any]     ,       __lowerCamelCase						:		List[str]     ,       __lowerCamelCase						:		List[str]     ,       __lowerCamelCase						:		int     ,       __lowerCamelCase						:		float = 1.0     ,       __lowerCamelCase						:		bool = True     ,       __lowerCamelCase						:		bool = False     ,       )		->							List[str]:
       SCREAMING_SNAKE_CASE__   =   self.controlnet_conditioning_channel_order
       if channel_order == "bgr":
            SCREAMING_SNAKE_CASE__   =   jnp.flip(__lowerCamelCase     ,       axis=1 )

       # 1. time
       if not isinstance(__lowerCamelCase     ,       jnp.ndarray ):
            SCREAMING_SNAKE_CASE__   =   jnp.array([timesteps]     ,       dtype=jnp.intaa )
       elif isinstance(__lowerCamelCase     ,       jnp.ndarray ) and len(timesteps.shape ) == 0:
            SCREAMING_SNAKE_CASE__   =   timesteps.astype(dtype=jnp.floataa )
            SCREAMING_SNAKE_CASE__   =   jnp.expand_dims(__lowerCamelCase     ,       0 )

       SCREAMING_SNAKE_CASE__   =   self.time_proj(__lowerCamelCase )
       SCREAMING_SNAKE_CASE__   =   self.time_embedding(__lowerCamelCase )

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

       SCREAMING_SNAKE_CASE__   =   jnp.transpose(__lowerCamelCase     ,       (0, 2, 3, 1) )
       SCREAMING_SNAKE_CASE__   =   self.controlnet_cond_embedding(__lowerCamelCase )
       sample += controlnet_cond

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

       # 4. mid
       SCREAMING_SNAKE_CASE__   =   self.mid_block(__lowerCamelCase     ,       __lowerCamelCase     ,       __lowerCamelCase     ,       deterministic=not train )

       # 5. contronet blocks
       SCREAMING_SNAKE_CASE__   =   ()
       for down_block_res_sample, controlnet_block in zip(__lowerCamelCase     ,       self.controlnet_down_blocks ):
            SCREAMING_SNAKE_CASE__   =   controlnet_block(__lowerCamelCase )
            controlnet_down_block_res_samples += (down_block_res_sample,)

       SCREAMING_SNAKE_CASE__   =   controlnet_down_block_res_samples

       SCREAMING_SNAKE_CASE__   =   self.controlnet_mid_block(__lowerCamelCase )

       # 6. scaling
       SCREAMING_SNAKE_CASE__   =   [sample * conditioning_scale for sample in down_block_res_samples]
       mid_block_res_sample *= conditioning_scale

       if not return_dict:
            return (down_block_res_samples, mid_block_res_sample)

       return FlaxControlNetOutput(
           down_block_res_samples=__lowerCamelCase     ,       mid_block_res_sample=__lowerCamelCase )





import datasets

from .evaluate import evaluate


_SCREAMING_SNAKE_CASE					:  Union[str, Any]             =      '''\
@inproceedings{Rajpurkar2016SQuAD10,
  title={SQuAD: 100, 000+ Questions for Machine Comprehension of Text},
  author={Pranav Rajpurkar and Jian Zhang and Konstantin Lopyrev and Percy Liang},
  booktitle={EMNLP},
  year={2016}
}
'''

_SCREAMING_SNAKE_CASE					:  Dict             =      '''
This metric wrap the official scoring script for version 1 of the Stanford Question Answering Dataset (SQuAD).

Stanford Question Answering Dataset (SQuAD) is a reading comprehension dataset, consisting of questions posed by
crowdworkers on a set of Wikipedia articles, where the answer to every question is a segment of text, or span,
from the corresponding reading passage, or the question might be unanswerable.
'''

_SCREAMING_SNAKE_CASE					:  str             =      '''
Computes SQuAD scores (F1 and EM).
Args:
    predictions: List of question-answers dictionaries with the following key-values:
        - \'id\': id of the question-answer pair as given in the references (see below)
        - \'prediction_text\': the text of the answer
    references: List of question-answers dictionaries with the following key-values:
        - \'id\': id of the question-answer pair (see above),
        - \'answers\': a Dict in the SQuAD dataset format
            {
                \'text\': list of possible texts for the answer, as a list of strings
                \'answer_start\': list of start positions for the answer, as a list of ints
            }
            Note that answer_start values are not taken into account to compute the metric.
Returns:
    \'exact_match\': Exact match (the normalized answer exactly match the gold answer)
    \'f1\': The F-score of predicted tokens versus the gold answer
Examples:

    >>> predictions = [{\'prediction_text\': \'1976\', \'id\': \'56e10a3be3433e1400422b22\'}]
    >>> references = [{\'answers\': {\'answer_start\': [97], \'text\': [\'1976\']}, \'id\': \'56e10a3be3433e1400422b22\'}]
    >>> squad_metric = datasets.load_metric("squad")
    >>> results = squad_metric.compute(predictions=predictions, references=references)
    >>> print(results)
    {\'exact_match\': 100.0, \'f1\': 100.0}
'''






@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION						,					_KWARGS_DESCRIPTION	)
class 	UpperCAmelCase__		(	datasets.Metric	):






  """simple docstring"""


  def 				lowercase_		(							self						:		List[Any] )		->							Optional[int]:
       return datasets.MetricInfo(
           description=_DESCRIPTION     ,       citation=_CITATION     ,       inputs_description=_KWARGS_DESCRIPTION     ,       features=datasets.Features(
               {
                   '''predictions''': {'''id''': datasets.Value('''string''' ), '''prediction_text''': datasets.Value('''string''' )},
                   '''references''': {
                       '''id''': datasets.Value('''string''' ),
                       '''answers''': datasets.features.Sequence(
                           {
                               '''text''': datasets.Value('''string''' ),
                               '''answer_start''': datasets.Value('''int32''' ),
                           } ),
                   },
               } )     ,       codebase_urls=['''https://rajpurkar.github.io/SQuAD-explorer/''']     ,       reference_urls=['''https://rajpurkar.github.io/SQuAD-explorer/''']     ,       )







  def 				lowercase_		(							self						:		Optional[Any]     ,       __lowerCamelCase						:		Dict     ,       __lowerCamelCase						:		List[Any] )		->							Optional[int]:
       SCREAMING_SNAKE_CASE__   =   {prediction['''id''']: prediction['''prediction_text'''] for prediction in predictions}
       SCREAMING_SNAKE_CASE__   =   [
           {
               '''paragraphs''': [
                   {
                       '''qas''': [
                           {
                               '''answers''': [{'''text''': answer_text} for answer_text in ref['''answers''']['''text''']],
                               '''id''': ref['''id'''],
                           }
                           for ref in references
                       ]
                   }
               ]
           }
       ]
       SCREAMING_SNAKE_CASE__   =   evaluate(dataset=__lowerCamelCase     ,       predictions=__lowerCamelCase )
       return score


