import json
import os
import unittest

from transformers.models.roc_bert.tokenization_roc_bert import (
    VOCAB_FILES_NAMES,
    RoCBertBasicTokenizer,
    RoCBertTokenizer,
    RoCBertWordpieceTokenizer,
    _is_control,
    _is_punctuation,
    _is_whitespace,
)
from transformers.testing_utils import require_tokenizers, slow

from ...test_tokenization_common import TokenizerTesterMixin, filter_non_english






@require_tokenizers
class 		SCREAMING_SNAKE_CASE__		(   _a    ,     unittest.TestCase    ):
    __a  =      RoCBertTokenizer
    __a  =      None
    __a  =      False
    __a  =      True
    __a  =      filter_non_english
    def 						UpperCamelCase_				(  self   )       ->					Any:
     super().setUp()

     SCREAMING_SNAKE_CASE__:   Union[str, Any]=      ['''[UNK]''', '''[CLS]''', '''[SEP]''', '''[PAD]''', '''[MASK]''', '''你''', '''好''', '''是''', '''谁''', '''a''', '''b''', '''c''', '''d''']
     SCREAMING_SNAKE_CASE__:   Dict=      {}
     SCREAMING_SNAKE_CASE__:   Optional[Any]=      {}
     for i, value in enumerate(_A   ):
      SCREAMING_SNAKE_CASE__:   Optional[Any]=      i
      SCREAMING_SNAKE_CASE__:   Any=      i
     SCREAMING_SNAKE_CASE__:   Optional[Any]=      os.path.join(self.tmpdirname , VOCAB_FILES_NAMES['''vocab_file''']   )
     SCREAMING_SNAKE_CASE__:   List[Any]=      os.path.join(self.tmpdirname , VOCAB_FILES_NAMES['''word_shape_file''']   )
     SCREAMING_SNAKE_CASE__:   Optional[int]=      os.path.join(self.tmpdirname , VOCAB_FILES_NAMES['''word_pronunciation_file''']   )
     with open(self.vocab_file , '''w''' , encoding='''utf-8'''   ) as vocab_writer:
      vocab_writer.write(''''''.join([x + '''\n''' for x in vocab_tokens]   )   )
     with open(self.word_shape_file , '''w''' , encoding='''utf-8'''   ) as word_shape_writer:
      json.dump(_A , _A , ensure_ascii=_A   )
     with open(self.word_pronunciation_file , '''w''' , encoding='''utf-8'''   ) as word_pronunciation_writer:
      json.dump(_A , _A , ensure_ascii=_A   )
    def 						UpperCamelCase_				(  self   )       ->					str:
     SCREAMING_SNAKE_CASE__:   Tuple=      self.tokenizer_class(self.vocab_file , self.word_shape_file , self.word_pronunciation_file   )

     SCREAMING_SNAKE_CASE__:   Dict=      tokenizer.tokenize('''你好[SEP]你是谁'''   )
     self.assertListEqual(_A , ['''你''', '''好''', '''[SEP]''', '''你''', '''是''', '''谁''']   )
     self.assertListEqual(tokenizer.convert_tokens_to_ids(_A   ) , [5, 6, 2, 5, 7, 8]   )
     self.assertListEqual(tokenizer.convert_tokens_to_shape_ids(_A   ) , [5, 6, 2, 5, 7, 8]   )
     self.assertListEqual(tokenizer.convert_tokens_to_pronunciation_ids(_A   ) , [5, 6, 2, 5, 7, 8]   )
    def 						UpperCamelCase_				(  self   )       ->					Tuple:
     SCREAMING_SNAKE_CASE__:   Optional[int]=      RoCBertBasicTokenizer()

     self.assertListEqual(tokenizer.tokenize('''ah\u535A\u63A8zz'''   ) , ['''ah''', '''\u535A''', '''\u63A8''', '''zz''']   )
    def 						UpperCamelCase_				(  self   )       ->					Union[str, Any]:
     SCREAMING_SNAKE_CASE__:   str=      RoCBertBasicTokenizer(do_lower_case=_A   )

     self.assertListEqual(
         tokenizer.tokenize(''' \tHeLLo!how  \n Are yoU?  '''   ) , ['''hello''', '''!''', '''how''', '''are''', '''you''', '''?''']   )
     self.assertListEqual(tokenizer.tokenize('''H\u00E9llo'''   ) , ['''hello''']   )
    def 						UpperCamelCase_				(  self   )       ->					Tuple:
     SCREAMING_SNAKE_CASE__:   int=      RoCBertBasicTokenizer(do_lower_case=_A , strip_accents=_A   )

     self.assertListEqual(
         tokenizer.tokenize(''' \tHäLLo!how  \n Are yoU?  '''   ) , ['''hällo''', '''!''', '''how''', '''are''', '''you''', '''?''']   )
     self.assertListEqual(tokenizer.tokenize('''H\u00E9llo'''   ) , ['''h\u00E9llo''']   )
    def 						UpperCamelCase_				(  self   )       ->					str:
     SCREAMING_SNAKE_CASE__:   Any=      RoCBertBasicTokenizer(do_lower_case=_A , strip_accents=_A   )

     self.assertListEqual(
         tokenizer.tokenize(''' \tHäLLo!how  \n Are yoU?  '''   ) , ['''hallo''', '''!''', '''how''', '''are''', '''you''', '''?''']   )
     self.assertListEqual(tokenizer.tokenize('''H\u00E9llo'''   ) , ['''hello''']   )
    def 						UpperCamelCase_				(  self   )       ->					List[str]:
     SCREAMING_SNAKE_CASE__:   int=      RoCBertBasicTokenizer(do_lower_case=_A   )

     self.assertListEqual(
         tokenizer.tokenize(''' \tHäLLo!how  \n Are yoU?  '''   ) , ['''hallo''', '''!''', '''how''', '''are''', '''you''', '''?''']   )
     self.assertListEqual(tokenizer.tokenize('''H\u00E9llo'''   ) , ['''hello''']   )
    def 						UpperCamelCase_				(  self   )       ->					str:
     SCREAMING_SNAKE_CASE__:   Dict=      RoCBertBasicTokenizer(do_lower_case=_A   )

     self.assertListEqual(
         tokenizer.tokenize(''' \tHeLLo!how  \n Are yoU?  '''   ) , ['''HeLLo''', '''!''', '''how''', '''Are''', '''yoU''', '''?''']   )
    def 						UpperCamelCase_				(  self   )       ->					Any:
     SCREAMING_SNAKE_CASE__:   Tuple=      RoCBertBasicTokenizer(do_lower_case=_A , strip_accents=_A   )

     self.assertListEqual(
         tokenizer.tokenize(''' \tHäLLo!how  \n Are yoU?  '''   ) , ['''HäLLo''', '''!''', '''how''', '''Are''', '''yoU''', '''?''']   )
    def 						UpperCamelCase_				(  self   )       ->					Any:
     SCREAMING_SNAKE_CASE__:   Any=      RoCBertBasicTokenizer(do_lower_case=_A , strip_accents=_A   )

     self.assertListEqual(
         tokenizer.tokenize(''' \tHäLLo!how  \n Are yoU?  '''   ) , ['''HaLLo''', '''!''', '''how''', '''Are''', '''yoU''', '''?''']   )
    def 						UpperCamelCase_				(  self   )       ->					Tuple:
     SCREAMING_SNAKE_CASE__:   Optional[int]=      RoCBertBasicTokenizer(do_lower_case=_A , never_split=['''[UNK]''']   )

     self.assertListEqual(
         tokenizer.tokenize(''' \tHeLLo!how  \n Are yoU? [UNK]'''   ) , ['''HeLLo''', '''!''', '''how''', '''Are''', '''yoU''', '''?''', '''[UNK]''']   )
    def 						UpperCamelCase_				(  self   )       ->					str:
     SCREAMING_SNAKE_CASE__:   List[Any]=      ['''[UNK]''', '''[CLS]''', '''[SEP]''', '''want''', '''##want''', '''##ed''', '''wa''', '''un''', '''runn''', '''##ing''']

     SCREAMING_SNAKE_CASE__:   Any=      {}
     for i, token in enumerate(_A   ):
      SCREAMING_SNAKE_CASE__:   Tuple=      i
     SCREAMING_SNAKE_CASE__:   Dict=      RoCBertWordpieceTokenizer(vocab=_A , unk_token='''[UNK]'''   )

     self.assertListEqual(tokenizer.tokenize(''''''   ) , []   )

     self.assertListEqual(tokenizer.tokenize('''unwanted running'''   ) , ['''un''', '''##want''', '''##ed''', '''runn''', '''##ing''']   )

     self.assertListEqual(tokenizer.tokenize('''unwantedX running'''   ) , ['''[UNK]''', '''runn''', '''##ing''']   )
    def 						UpperCamelCase_				(  self   )       ->					Tuple:
     self.assertTrue(_is_whitespace(''' '''   )   )
     self.assertTrue(_is_whitespace('''\t'''   )   )
     self.assertTrue(_is_whitespace('''\r'''   )   )
     self.assertTrue(_is_whitespace('''\n'''   )   )
     self.assertTrue(_is_whitespace('''\u00A0'''   )   )

     self.assertFalse(_is_whitespace('''A'''   )   )
     self.assertFalse(_is_whitespace('''-'''   )   )
    def 						UpperCamelCase_				(  self   )       ->					Any:
     self.assertTrue(_is_control('''\u0005'''   )   )

     self.assertFalse(_is_control('''A'''   )   )
     self.assertFalse(_is_control(''' '''   )   )
     self.assertFalse(_is_control('''\t'''   )   )
     self.assertFalse(_is_control('''\r'''   )   )
    def 						UpperCamelCase_				(  self   )       ->					Dict:
     self.assertTrue(_is_punctuation('''-'''   )   )
     self.assertTrue(_is_punctuation('''$'''   )   )
     self.assertTrue(_is_punctuation('''`'''   )   )
     self.assertTrue(_is_punctuation('''.'''   )   )

     self.assertFalse(_is_punctuation('''A'''   )   )
     self.assertFalse(_is_punctuation(''' '''   )   )
    def 						UpperCamelCase_				(  self   )       ->					int:
     SCREAMING_SNAKE_CASE__:   List[Any]=      self.get_tokenizer()

     # Example taken from the issue https://github.com/huggingface/tokenizers/issues/340
     self.assertListEqual([tokenizer.tokenize(_A   ) for t in ['''Test''', '''\xad''', '''test''']] , [['''[UNK]'''], [], ['''[UNK]''']]   )

     if self.test_rust_tokenizer:
      SCREAMING_SNAKE_CASE__:   Union[str, Any]=      self.get_rust_tokenizer()
      self.assertListEqual(
          [rust_tokenizer.tokenize(_A   ) for t in ['''Test''', '''\xad''', '''test''']] , [['''[UNK]'''], [], ['''[UNK]''']]   )
    def 						UpperCamelCase_				(  self   )       ->					Optional[int]:
     for tokenizer, pretrained_name, kwargs in self.tokenizers_list:
      with self.subTest(f'{tokenizer.__class__.__name__} ({pretrained_name})'   ):
       SCREAMING_SNAKE_CASE__:   Dict=      self.rust_tokenizer_class.from_pretrained(_A , **_A   )

       SCREAMING_SNAKE_CASE__:   List[Any]=      f'A, naïve {tokenizer_r.mask_token} AllenNLP sentence.'
       SCREAMING_SNAKE_CASE__:   Union[str, Any]=      tokenizer_r.encode_plus(
           _A , return_attention_mask=_A , return_token_type_ids=_A , return_offsets_mapping=_A , add_special_tokens=_A , )

       SCREAMING_SNAKE_CASE__:   Any=      tokenizer_r.do_lower_case if hasattr(_A , '''do_lower_case'''   ) else False
       SCREAMING_SNAKE_CASE__:   List[str]=      (
           [
               ((0, 0), tokenizer_r.cls_token),
               ((0, 1), '''A'''),
               ((1, 2), ''','''),
               ((3, 5), '''na'''),
               ((5, 6), '''##ï'''),
               ((6, 8), '''##ve'''),
               ((9, 15), tokenizer_r.mask_token),
               ((16, 21), '''Allen'''),
               ((21, 23), '''##NL'''),
               ((23, 24), '''##P'''),
               ((25, 33), '''sentence'''),
               ((33, 34), '''.'''),
               ((0, 0), tokenizer_r.sep_token),
           ]
           if not do_lower_case
           else [
               ((0, 0), tokenizer_r.cls_token),
               ((0, 1), '''a'''),
               ((1, 2), ''','''),
               ((3, 8), '''naive'''),
               ((9, 15), tokenizer_r.mask_token),
               ((16, 21), '''allen'''),
               ((21, 23), '''##nl'''),
               ((23, 24), '''##p'''),
               ((25, 33), '''sentence'''),
               ((33, 34), '''.'''),
               ((0, 0), tokenizer_r.sep_token),
           ]
       )

       self.assertEqual(
           [e[1] for e in expected_results] , tokenizer_r.convert_ids_to_tokens(tokens['''input_ids''']   )   )
       self.assertEqual([e[0] for e in expected_results] , tokens['''offset_mapping''']   )
    def 						UpperCamelCase_				(  self   )       ->					Union[str, Any]:
     SCREAMING_SNAKE_CASE__:   Dict=      ['''的''', '''人''', '''有''']
     SCREAMING_SNAKE_CASE__:   str=      ''''''.join(_A   )
     for tokenizer, pretrained_name, kwargs in self.tokenizers_list:
      with self.subTest(f'{tokenizer.__class__.__name__} ({pretrained_name})'   ):
       SCREAMING_SNAKE_CASE__:   Dict=      True
       SCREAMING_SNAKE_CASE__:   str=      self.tokenizer_class.from_pretrained(_A , **_A   )
       SCREAMING_SNAKE_CASE__:   Dict=      self.rust_tokenizer_class.from_pretrained(_A , **_A   )

       SCREAMING_SNAKE_CASE__:   List[Any]=      tokenizer_p.encode(_A , add_special_tokens=_A   )
       SCREAMING_SNAKE_CASE__:   Dict=      tokenizer_r.encode(_A , add_special_tokens=_A   )

       SCREAMING_SNAKE_CASE__:   List[str]=      tokenizer_r.convert_ids_to_tokens(_A   )
       SCREAMING_SNAKE_CASE__:   Dict=      tokenizer_p.convert_ids_to_tokens(_A   )

       # it is expected that each Chinese character is not preceded by "##"
       self.assertListEqual(_A , _A   )
       self.assertListEqual(_A , _A   )

       SCREAMING_SNAKE_CASE__:   List[str]=      False
       SCREAMING_SNAKE_CASE__:   str=      self.rust_tokenizer_class.from_pretrained(_A , **_A   )
       SCREAMING_SNAKE_CASE__:   Any=      self.tokenizer_class.from_pretrained(_A , **_A   )

       SCREAMING_SNAKE_CASE__:   Tuple=      tokenizer_r.encode(_A , add_special_tokens=_A   )
       SCREAMING_SNAKE_CASE__:   Tuple=      tokenizer_p.encode(_A , add_special_tokens=_A   )

       SCREAMING_SNAKE_CASE__:   Dict=      tokenizer_r.convert_ids_to_tokens(_A   )
       SCREAMING_SNAKE_CASE__:   Optional[int]=      tokenizer_p.convert_ids_to_tokens(_A   )

       # it is expected that only the first Chinese character is not preceded by "##".
       SCREAMING_SNAKE_CASE__:   Optional[int]=      [
           f'##{token}' if idx != 0 else token for idx, token in enumerate(_A   )
       ]
       self.assertListEqual(_A , _A   )
       self.assertListEqual(_A , _A   )
    @slow
    def 						UpperCamelCase_				(  self   )       ->					Tuple:
     SCREAMING_SNAKE_CASE__:   Optional[int]=      self.tokenizer_class(self.vocab_file , self.word_shape_file , self.word_pronunciation_file   )

     SCREAMING_SNAKE_CASE__:   List[Any]=      tokenizer.encode('''你好''' , add_special_tokens=_A   )
     SCREAMING_SNAKE_CASE__:   int=      tokenizer.encode('''你是谁''' , add_special_tokens=_A   )

     SCREAMING_SNAKE_CASE__:   List[str]=      tokenizer.build_inputs_with_special_tokens(_A   )
     SCREAMING_SNAKE_CASE__:   Union[str, Any]=      tokenizer.build_inputs_with_special_tokens(_A , _A   )

     assert encoded_sentence == [1] + text + [2]
     assert encoded_pair == [1] + text + [2] + text_a + [2]



    def 						UpperCamelCase_				(  self   )       ->					str:
     SCREAMING_SNAKE_CASE__:   Any=      self.get_tokenizers(do_lower_case=_A   )
     for tokenizer in tokenizers:
      with self.subTest(f'{tokenizer.__class__.__name__}'   ):
       SCREAMING_SNAKE_CASE__:   Tuple=      '''你好，你是谁'''
       SCREAMING_SNAKE_CASE__:   str=      tokenizer.tokenize(_A   )
       SCREAMING_SNAKE_CASE__:   Dict=      tokenizer.convert_tokens_to_ids(_A   )
       SCREAMING_SNAKE_CASE__:   List[str]=      tokenizer.convert_tokens_to_shape_ids(_A   )
       SCREAMING_SNAKE_CASE__:   Tuple=      tokenizer.convert_tokens_to_pronunciation_ids(_A   )
       SCREAMING_SNAKE_CASE__:   Dict=      tokenizer.prepare_for_model(
           _A , _A , _A , add_special_tokens=_A   )

       SCREAMING_SNAKE_CASE__:   Dict=      tokenizer.encode_plus(_A , add_special_tokens=_A   )

       self.assertEqual(_A , _A   )


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 		_lowerCamelCase		(   UpperCamelCase_    ):
    __a  =      42
    __a  =      42






class 		_lowerCamelCase		(   nn.Module    ):
    __a  =      42
    __a  =      (16, 32, 96, 256)
    __a  =      jnp.floataa
    def 						UpperCamelCase_				(  self   )       ->					Optional[int]:
     SCREAMING_SNAKE_CASE__:   List[str]=      nn.Conv(
         self.block_out_channels[0] , kernel_size=(3, 3) , padding=((1, 1), (1, 1)) , dtype=self.dtype , )

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

     SCREAMING_SNAKE_CASE__:   Optional[Any]=      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 , lowerCAmelCase   )       ->					Optional[Any]:
     SCREAMING_SNAKE_CASE__:   Dict=      self.conv_in(lowerCAmelCase   )
     SCREAMING_SNAKE_CASE__:   Optional[int]=      nn.silu(lowerCAmelCase   )

     for block in self.blocks:
      SCREAMING_SNAKE_CASE__:   Optional[int]=      block(lowerCAmelCase   )
      SCREAMING_SNAKE_CASE__:   Optional[Any]=      nn.silu(lowerCAmelCase   )

     SCREAMING_SNAKE_CASE__:   Any=      self.conv_out(lowerCAmelCase   )

     return embedding






@flax_register_to_config
class 		_lowerCamelCase		(   nn.Module    ,     UpperCamelCase_    ,     UpperCamelCase_    ):
    __a  =      32
    __a  =      4
    __a  =      (
        "CrossAttnDownBlock2D",
        "CrossAttnDownBlock2D",
        "CrossAttnDownBlock2D",
        "DownBlock2D",
    )
    __a  =      False
    __a  =      (320, 640, 1280, 1280)
    __a  =      2
    __a  =      8
    __a  =      None
    __a  =      1280
    __a  =      0.0
    __a  =      False
    __a  =      jnp.floataa
    __a  =      True
    __a  =      0
    __a  =      "rgb"
    __a  =      (16, 32, 96, 256)
    def 						UpperCamelCase_				(  self , lowerCAmelCase   )       ->					FrozenDict:
     # init input tensors
     SCREAMING_SNAKE_CASE__:   Any=      (1, self.in_channels, self.sample_size, self.sample_size)
     SCREAMING_SNAKE_CASE__:   Union[str, Any]=      jnp.zeros(lowerCAmelCase , dtype=jnp.floataa   )
     SCREAMING_SNAKE_CASE__:   Union[str, Any]=      jnp.ones((1,) , dtype=jnp.intaa   )
     SCREAMING_SNAKE_CASE__:   Optional[Any]=      jnp.zeros((1, 1, self.cross_attention_dim) , dtype=jnp.floataa   )
     SCREAMING_SNAKE_CASE__:   int=      (1, 3, self.sample_size * 8, self.sample_size * 8)
     SCREAMING_SNAKE_CASE__:   Optional[int]=      jnp.zeros(lowerCAmelCase , dtype=jnp.floataa   )

     SCREAMING_SNAKE_CASE__,	SCREAMING_SNAKE_CASE__:   Any=      jax.random.split(lowerCAmelCase   )
     SCREAMING_SNAKE_CASE__:   Optional[int]=      {'''params''': params_rng, '''dropout''': dropout_rng}

     return self.init(lowerCAmelCase , lowerCAmelCase , lowerCAmelCase , lowerCAmelCase , lowerCAmelCase   )["params"]
    def 						UpperCamelCase_				(  self   )       ->					Optional[Any]:
     SCREAMING_SNAKE_CASE__:   Dict=      self.block_out_channels
     SCREAMING_SNAKE_CASE__:   str=      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__:   int=      self.num_attention_heads or self.attention_head_dim

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

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

     SCREAMING_SNAKE_CASE__:   List[str]=      self.only_cross_attention
     if isinstance(lowerCAmelCase , lowerCAmelCase   ):
      SCREAMING_SNAKE_CASE__:   str=      (only_cross_attention,) * len(self.down_block_types   )

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

     # down
     SCREAMING_SNAKE_CASE__:   Optional[int]=      []
     SCREAMING_SNAKE_CASE__:   List[str]=      []

     SCREAMING_SNAKE_CASE__:   Union[str, Any]=      block_out_channels[0]

     SCREAMING_SNAKE_CASE__:   Dict=      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__:   Tuple=      output_channel
      SCREAMING_SNAKE_CASE__:   Any=      block_out_channels[i]
      SCREAMING_SNAKE_CASE__:   Optional[Any]=      i == len(lowerCAmelCase   ) - 1

      if down_block_type == "CrossAttnDownBlock2D":
       SCREAMING_SNAKE_CASE__:   Any=      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__:   Union[str, Any]=      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__:   str=      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__:   str=      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__:   int=      down_blocks
     SCREAMING_SNAKE_CASE__:   int=      controlnet_down_blocks

     # mid
     SCREAMING_SNAKE_CASE__:   Optional[int]=      block_out_channels[-1]
     SCREAMING_SNAKE_CASE__:   int=      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__:   str=      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 , lowerCAmelCase , lowerCAmelCase , lowerCAmelCase , lowerCAmelCase , lowerCAmelCase = 1.0 , lowerCAmelCase = True , lowerCAmelCase = False , )       ->					Union[FlaxControlNetOutput, Tuple]:
     SCREAMING_SNAKE_CASE__:   Dict=      self.controlnet_conditioning_channel_order
     if channel_order == "bgr":
      SCREAMING_SNAKE_CASE__:   Optional[Any]=      jnp.flip(lowerCAmelCase , axis=1   )

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

     SCREAMING_SNAKE_CASE__:   Dict=      self.time_proj(lowerCAmelCase   )
     SCREAMING_SNAKE_CASE__:   Dict=      self.time_embedding(lowerCAmelCase   )

     # 2. pre-process
     SCREAMING_SNAKE_CASE__:   int=      jnp.transpose(lowerCAmelCase , (0, 2, 3, 1)   )
     SCREAMING_SNAKE_CASE__:   str=      self.conv_in(lowerCAmelCase   )

     SCREAMING_SNAKE_CASE__:   Optional[Any]=      jnp.transpose(lowerCAmelCase , (0, 2, 3, 1)   )
     SCREAMING_SNAKE_CASE__:   str=      self.controlnet_cond_embedding(lowerCAmelCase   )
     sample += controlnet_cond

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

     # 4. mid
     SCREAMING_SNAKE_CASE__:   Optional[Any]=      self.mid_block(lowerCAmelCase , lowerCAmelCase , lowerCAmelCase , deterministic=not train   )

     # 5. contronet blocks
     SCREAMING_SNAKE_CASE__:   Optional[int]=      ()
     for down_block_res_sample, controlnet_block in zip(lowerCAmelCase , self.controlnet_down_blocks   ):
      SCREAMING_SNAKE_CASE__:   List[str]=      controlnet_block(lowerCAmelCase   )
      controlnet_down_block_res_samples += (down_block_res_sample,)

     SCREAMING_SNAKE_CASE__:   Optional[Any]=      controlnet_down_block_res_samples

     SCREAMING_SNAKE_CASE__:   Optional[int]=      self.controlnet_mid_block(lowerCAmelCase   )

     # 6. scaling
     SCREAMING_SNAKE_CASE__:   Tuple=      [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   )

