






"""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    (  _a  ):
   a      :				jnp.ndarray







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

   def 						SCREAMING_SNAKE_CASE__ (					self ,				snake_case_      )		->       FrozenDict:
       # init input tensors
       UpperCamelCase__         =    (1, self.in_channels, self.sample_size, self.sample_size)
       UpperCamelCase__         =    jnp.zeros(snake_case_ ,				dtype=jnp.floataa      )
       UpperCamelCase__         =    jnp.ones((1,) ,				dtype=jnp.intaa      )
       UpperCamelCase__         =    jnp.zeros((1, 1, self.cross_attention_dim) ,				dtype=jnp.floataa      )

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

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

   def 						SCREAMING_SNAKE_CASE__ (					self      )		->       List[Any]:
       UpperCamelCase__         =    self.block_out_channels
       UpperCamelCase__         =    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.
       UpperCamelCase__         =    self.num_attention_heads or self.attention_head_dim

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

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

       UpperCamelCase__         =    self.only_cross_attention
       if isinstance(snake_case_ ,				snake_case_      ):
           UpperCamelCase__         =    (only_cross_attention,) * len(self.down_block_types      )

       if isinstance(snake_case_ ,				snake_case_      ):
           UpperCamelCase__         =    (num_attention_heads,) * len(self.down_block_types      )

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

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

           down_blocks.append(snake_case_      )
       UpperCamelCase__         =    down_blocks

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

           UpperCamelCase__         =    i == len(snake_case_      ) - 1

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

           up_blocks.append(snake_case_      )
           UpperCamelCase__         =    output_channel
       UpperCamelCase__         =    up_blocks

       # out
       UpperCamelCase__         =    nn.GroupNorm(num_groups=32 ,				epsilon=1E-5      )
       UpperCamelCase__         =    nn.Conv(
           self.out_channels ,				kernel_size=(3, 3) ,				strides=(1, 1) ,				padding=((1, 1), (1, 1)) ,				dtype=self.dtype ,				)







   def __call__(					self ,				snake_case_ ,				snake_case_ ,				snake_case_ ,				snake_case_=None ,				snake_case_=None ,				snake_case_ = True ,				snake_case_ = False ,				)		->       Union[FlaxUNetaDConditionOutput, Tuple]:
       # 1. time
       if not isinstance(snake_case_ ,				jnp.ndarray      ):
           UpperCamelCase__         =    jnp.array([timesteps] ,				dtype=jnp.intaa      )
       elif isinstance(snake_case_ ,				jnp.ndarray      ) and len(timesteps.shape      ) == 0:
           UpperCamelCase__         =    timesteps.astype(dtype=jnp.floataa      )
           UpperCamelCase__         =    jnp.expand_dims(snake_case_ ,				0      )

       UpperCamelCase__         =    self.time_proj(snake_case_      )
       UpperCamelCase__         =    self.time_embedding(snake_case_      )

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

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

       if down_block_additional_residuals is not None:
           UpperCamelCase__         =    ()

           for down_block_res_sample, down_block_additional_residual in zip(
               snake_case_ ,				snake_case_      ):
               down_block_res_sample += down_block_additional_residual
               new_down_block_res_samples += (down_block_res_sample,)

           UpperCamelCase__         =    new_down_block_res_samples

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

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

       if not return_dict:
           return (sample,)

       return FlaxUNetaDConditionOutput(sample=snake_case_      )












"""simple docstring"""





import doctest
import logging
import os
import unittest
from pathlib import Path
from typing import List, Union

import transformers
from transformers.testing_utils import require_tf, require_torch, slow


A__					: Union[str, Any]=	logging.getLogger()
@unittest.skip("""Temporarily disable the doc tests."""  )
@require_torch
@require_tf
@slow
class    __lowerCamelCase    (  unittest.TestCase  ):

   def 						SCREAMING_SNAKE_CASE__ (					self ,				snake_case_ ,				snake_case_ = None ,				snake_case_ = None ,				snake_case_ = None ,				snake_case_ = True ,				)		->       Tuple:
       UpperCamelCase__         =    [file for file in os.listdir(snake_case_      ) if os.path.isfile(os.path.join(snake_case_ ,				snake_case_      )      )]

       if identifier is not None:
           UpperCamelCase__         =    [file for file in files if identifier in file]

       if n_identifier is not None:
           if isinstance(snake_case_ ,				snake_case_      ):
               for n_ in n_identifier:
                   UpperCamelCase__         =    [file for file in files if n_ not in file]
           else:
               UpperCamelCase__         =    [file for file in files if n_identifier not in file]

       UpperCamelCase__         =    ignore_files or []
       ignore_files.append('__init__.py'      )
       UpperCamelCase__         =    [file for file in files if file not in ignore_files]

       for file in files:
           # Open all files
           print('Testing' ,				snake_case_      )

           if only_modules:
               UpperCamelCase__         =    file.split('.'      )[0]
               try:
                   UpperCamelCase__         =    getattr(snake_case_ ,				snake_case_      )
                   UpperCamelCase__         =    doctest.DocTestSuite(snake_case_      )
                   UpperCamelCase__         =    unittest.TextTestRunner().run(snake_case_      )
                   self.assertIs(len(result.failures      ) ,				0      )
               except AttributeError:
                   logger.info(F'{module_identifier} is not a module.'      )
           else:
               UpperCamelCase__         =    doctest.testfile(str('..' / directory / file      ) ,				optionflags=doctest.ELLIPSIS      )
               self.assertIs(result.failed ,				0      )

   def 						SCREAMING_SNAKE_CASE__ (					self      )		->       Union[str, Any]:
       UpperCamelCase__         =    Path('src/transformers'      )
       UpperCamelCase__         =    'modeling'
       UpperCamelCase__         =    [
           'modeling_ctrl.py',
           'modeling_tf_ctrl.py',
       ]
       self.analyze_directory(snake_case_ ,				identifier=snake_case_ ,				ignore_files=snake_case_      )

   def 						SCREAMING_SNAKE_CASE__ (					self      )		->       Optional[int]:
       UpperCamelCase__         =    Path('src/transformers'      )
       UpperCamelCase__         =    'tokenization'
       self.analyze_directory(snake_case_ ,				identifier=snake_case_      )

   def 						SCREAMING_SNAKE_CASE__ (					self      )		->       Optional[Any]:
       UpperCamelCase__         =    Path('src/transformers'      )
       UpperCamelCase__         =    'configuration'
       self.analyze_directory(snake_case_ ,				identifier=snake_case_      )

   def 						SCREAMING_SNAKE_CASE__ (					self      )		->       str:
       UpperCamelCase__         =    Path('src/transformers'      )
       UpperCamelCase__         =    ['configuration', 'modeling', 'tokenization']
       self.analyze_directory(snake_case_ ,				n_identifier=snake_case_      )







   def 						SCREAMING_SNAKE_CASE__ (					self      )		->       Union[str, Any]:
       UpperCamelCase__         =    Path('docs/source'      )
       UpperCamelCase__         =    ['favicon.ico']
       self.analyze_directory(snake_case_ ,				ignore_files=snake_case_ ,				only_modules=snake_case_      )




