"""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_      (     lowerCAmelCase   ):



    """simple docstring"""
    _lowerCAmelCase       :      jnp.ndarray





@flax_register_to_config
class 					lowerCAmelCase_      (     nn.Module       ,      lowerCAmelCase       ,      lowerCAmelCase   ):



    """simple docstring"""
    _lowerCAmelCase       :      int     					=  32
    _lowerCAmelCase       :      int     					=  4
    _lowerCAmelCase       :      int     					=  4
    _lowerCAmelCase       :      Tuple[str]     					=  (
        "CrossAttnDownBlock2D",
        "CrossAttnDownBlock2D",
        "CrossAttnDownBlock2D",
        "DownBlock2D",
    )
    _lowerCAmelCase       :      Tuple[str]     					=  ("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")
    _lowerCAmelCase       :      Union[bool, Tuple[bool]]     					=  False
    _lowerCAmelCase       :      Tuple[int]     					=  (320, 640, 1280, 1280)
    _lowerCAmelCase       :      int     					=  2
    _lowerCAmelCase       :      Union[int, Tuple[int]]     					=  8
    _lowerCAmelCase       :      Optional[Union[int, Tuple[int]]]     					=  None
    _lowerCAmelCase       :      int     					=  1280
    _lowerCAmelCase       :      float     					=  0.0
    _lowerCAmelCase       :      bool     					=  False
    _lowerCAmelCase       :      jnp.dtype     					=  jnp.floataa
    _lowerCAmelCase       :      bool     					=  True
    _lowerCAmelCase       :      int     					=  0
    _lowerCAmelCase       :      bool     					=  False

    def 						snake_case	(			self						,      lowerCAmelCase      ):






          """simple docstring"""

          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 						snake_case	(			self      ):






          """simple docstring"""

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






          """simple docstring"""

          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      )



"""simple docstring"""
import os
import sys
import unittest


SCREAMING_SNAKE_CASE__   =					os.path.abspath(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))
sys.path.append(os.path.join(git_repo_path, "utils"))

import check_dummies  # noqa: E402
from check_dummies import create_dummy_files, create_dummy_object, find_backend, read_init  # noqa: E402


# Align TRANSFORMERS_PATH in check_dummies with the current path
SCREAMING_SNAKE_CASE__   =					os.path.join(git_repo_path, "src", "transformers")

SCREAMING_SNAKE_CASE__   =					"\n{0} = None\n"

SCREAMING_SNAKE_CASE__   =					"\nclass {0}(metaclass=DummyObject):\n    _backends = {1}\n\n    def __init__(self, *args, **kwargs):\n        requires_backends(self, {1})\n"


SCREAMING_SNAKE_CASE__   =					"\ndef {0}(*args, **kwargs):\n    requires_backends({0}, {1})\n"



class 					lowerCAmelCase_      (     unittest.TestCase   ):



				"""simple docstring"""

				def 						snake_case	(			self      ):






										"""simple docstring"""

										snake_case        =     find_backend('    _import_structure["models.albert"].append("AlbertTokenizerFast")'      )
										self.assertIsNone(lowerCAmelCase      )

										snake_case        =     find_backend('    if not is_tokenizers_available():'      )
										self.assertEqual(lowerCAmelCase						,      'tokenizers'      )

										snake_case        =     find_backend('    if not is_tensorflow_text_available():'      )
										self.assertEqual(lowerCAmelCase						,      'tensorflow_text'      )

										snake_case        =     find_backend('    if not (is_sentencepiece_available() and is_tokenizers_available()):'      )
										self.assertEqual(lowerCAmelCase						,      'sentencepiece_and_tokenizers'      )

										snake_case        =     find_backend(
										    '    if not (is_sentencepiece_available() and is_tensorflow_text_available()):'      )
										self.assertEqual(lowerCAmelCase						,      'sentencepiece_and_tensorflow_text'      )

										snake_case        =     find_backend(
										    '    if not (is_sentencepiece_available() and is_tokenizers_available() and is_vision_available()):'      )
										self.assertEqual(lowerCAmelCase						,      'sentencepiece_and_tokenizers_and_vision'      )

				def 						snake_case	(			self      ):






										"""simple docstring"""

										snake_case        =     read_init()
										# We don't assert on the exact list of keys to allow for smooth grow of backend-specific objects
										self.assertIn('torch'						,      lowerCAmelCase      )
										self.assertIn('tensorflow_text'						,      lowerCAmelCase      )
										self.assertIn('sentencepiece_and_tokenizers'						,      lowerCAmelCase      )

										# Likewise, we can't assert on the exact content of a key
										self.assertIn('BertModel'						,      objects['torch']      )
										self.assertIn('TFBertModel'						,      objects['tf']      )
										self.assertIn('FlaxBertModel'						,      objects['flax']      )
										self.assertIn('BertModel'						,      objects['torch']      )
										self.assertIn('TFBertTokenizer'						,      objects['tensorflow_text']      )
										self.assertIn('convert_slow_tokenizer'						,      objects['sentencepiece_and_tokenizers']      )

				def 						snake_case	(			self      ):






										"""simple docstring"""

										snake_case        =     create_dummy_object('CONSTANT'						,      '\'torch\''      )
										self.assertEqual(lowerCAmelCase						,      '\nCONSTANT = None\n'      )

										snake_case        =     create_dummy_object('function'						,      '\'torch\''      )
										self.assertEqual(
										    lowerCAmelCase						,      '\ndef function(*args, **kwargs):\n    requires_backends(function, \'torch\')\n'      )

										snake_case        =     '\nclass FakeClass(metaclass=DummyObject):\n    _backends = \'torch\'\n\n    def __init__(self, *args, **kwargs):\n        requires_backends(self, \'torch\')\n'
										snake_case        =     create_dummy_object('FakeClass'						,      '\'torch\''      )
										self.assertEqual(lowerCAmelCase						,      lowerCAmelCase      )

				def 						snake_case	(			self      ):






										"""simple docstring"""

										snake_case        =     '# This file is autogenerated by the command `make fix-copies`, do not edit.\nfrom ..utils import DummyObject, requires_backends\n\n\nCONSTANT = None\n\n\ndef function(*args, **kwargs):\n    requires_backends(function, ["torch"])\n\n\nclass FakeClass(metaclass=DummyObject):\n    _backends = ["torch"]\n\n    def __init__(self, *args, **kwargs):\n        requires_backends(self, ["torch"])\n'
										snake_case        =     create_dummy_files({'torch': ['CONSTANT', 'function', 'FakeClass']}      )
										self.assertEqual(dummy_files['torch']						,      lowerCAmelCase      )


