



"""simple docstring"""
import os
import pickle
import unittest

from transformers import AutoTokenizer
from transformers.models.bert.tokenization_bert import BertTokenizer
from transformers.models.bert_japanese.tokenization_bert_japanese import (
    VOCAB_FILES_NAMES,
    BertJapaneseTokenizer,
    CharacterTokenizer,
    JumanppTokenizer,
    MecabTokenizer,
    SudachiTokenizer,
    WordpieceTokenizer,
)
from transformers.testing_utils import custom_tokenizers, require_jumanpp, require_sudachi

from ...test_tokenization_common import TokenizerTesterMixin




@custom_tokenizers
class 			snake_case_						(						a_     ,unittest.TestCase				):
							__lowerCAmelCase    					=					BertJapaneseTokenizer
							__lowerCAmelCase    					=					False
							__lowerCAmelCase    					=					True






							def snake_case_  (   self      ):
									super().setUp()

									a_     :      str 						=     [
									    "[UNK]",
									    "[CLS]",
									    "[SEP]",
									    "こんにちは",
									    "こん",
									    "にちは",
									    "ばんは",
									    "##こん",
									    "##にちは",
									    "##ばんは",
									    "世界",
									    "##世界",
									    "、",
									    "##、",
									    "。",
									    "##。",
									]

									a_     :      Tuple 						=     os.path.join(self.tmpdirname			,     VOCAB_FILES_NAMES["vocab_file"]      )
									with open(self.vocab_file			,     "w"			,     encoding="utf-8"      ) as vocab_writer:
											vocab_writer.write("".join([x + "\n" for x in vocab_tokens]      )      )






							def snake_case_  (   self			,     a_      ):
									a_     :      Dict 						=     "こんにちは、世界。 \nこんばんは、世界。"
									a_     :      int 						=     "こんにちは 、 世界 。 こんばんは 、 世界 。"
									return input_text, output_text






							def snake_case_  (   self			,     a_      ):
									a_						,				a_     :      List[str] 						=     self.get_input_output_texts(a_      )
									a_     :      Tuple 						=     tokenizer.encode(a_			,     add_special_tokens=a_      )
									a_     :      Dict 						=     tokenizer.decode(a_			,     clean_up_tokenization_spaces=a_      )
									return text, ids






							def snake_case_  (   self      ):
									pass  # TODO add if relevant






							def snake_case_  (   self      ):
									pass  # TODO add if relevant






							def snake_case_  (   self      ):
									pass  # TODO add if relevant






							def snake_case_  (   self      ):
									a_     :      int 						=     self.tokenizer_class(self.vocab_file      )

									a_     :      List[Any] 						=     tokenizer.tokenize("こんにちは、世界。\nこんばんは、世界。"      )
									self.assertListEqual(a_			,     ["こんにちは", "、", "世界", "。", "こん", "##ばんは", "、", "世界", "。"]      )
									self.assertListEqual(tokenizer.convert_tokens_to_ids(a_      )			,     [3, 1_2, 1_0, 1_4, 4, 9, 1_2, 1_0, 1_4]      )






							def snake_case_  (   self      ):
									a_     :      List[str] 						=     self.tokenizer_class(self.vocab_file			,     word_tokenizer_type="mecab"      )
									self.assertIsNotNone(a_      )

									a_     :      Union[str, Any] 						=     "こんにちは、世界。\nこんばんは、世界。"
									a_     :      Union[str, Any] 						=     tokenizer.tokenize(a_      )
									self.assertListEqual(a_			,     ["こんにちは", "、", "世界", "。", "こん", "##ばんは", "、", "世界", "。"]      )
									self.assertListEqual(tokenizer.convert_tokens_to_ids(a_      )			,     [3, 1_2, 1_0, 1_4, 4, 9, 1_2, 1_0, 1_4]      )

									a_     :      Any 						=     os.path.join(self.tmpdirname			,     "tokenizer.bin"      )
									with open(a_			,     "wb"      ) as handle:
											pickle.dump(a_			,     a_      )

									with open(a_			,     "rb"      ) as handle:
											a_     :      Optional[Any] 						=     pickle.load(a_      )

									a_     :      int 						=     tokenizer_new.tokenize(a_      )

									self.assertListEqual(a_			,     a_      )






							def snake_case_  (   self      ):
									a_     :      Tuple 						=     MecabTokenizer(mecab_dic="ipadic"      )

									self.assertListEqual(
									    tokenizer.tokenize(" \tｱｯﾌﾟﾙストアでiPhone８ が  \n 発売された　。  "      )			,     ["アップルストア", "で", "iPhone", "8", "が", "発売", "さ", "れ", "た", "。"]			,     )






							def snake_case_  (   self      ):
									try:
											a_     :      Dict 						=     MecabTokenizer(mecab_dic="unidic_lite"      )
									except ModuleNotFoundError:
											return

									self.assertListEqual(
									    tokenizer.tokenize(" \tｱｯﾌﾟﾙストアでiPhone８ が  \n 発売された　。  "      )			,     ["アップル", "ストア", "で", "iPhone", "8", "が", "発売", "さ", "れ", "た", "。"]			,     )






							def snake_case_  (   self      ):
									try:
											a_     :      List[Any] 						=     MecabTokenizer(mecab_dic="unidic"      )
									except ModuleNotFoundError:
											return

									self.assertListEqual(
									    tokenizer.tokenize(" \tｱｯﾌﾟﾙストアでiPhone８ が  \n 発売された　。  "      )			,     ["アップル", "ストア", "で", "iPhone", "8", "が", "発売", "さ", "れ", "た", "。"]			,     )






							def snake_case_  (   self      ):
									a_     :      List[str] 						=     MecabTokenizer(do_lower_case=a_			,     mecab_dic="ipadic"      )

									self.assertListEqual(
									    tokenizer.tokenize(" \tｱｯﾌﾟﾙストアでiPhone８ が  \n 発売された　。  "      )			,     ["アップルストア", "で", "iphone", "8", "が", "発売", "さ", "れ", "た", "。"]			,     )






							def snake_case_  (   self      ):
									try:
											a_     :      int 						=     MecabTokenizer(
											    do_lower_case=a_			,     normalize_text=a_			,     mecab_option="-d /usr/local/lib/mecab/dic/jumandic"      )
									except RuntimeError:
											# if dict doesn't exist in the system, previous code raises this error.
											return

									self.assertListEqual(
									    tokenizer.tokenize(" \tｱｯﾌﾟﾙストアでiPhone８ が  \n 発売された　。  "      )			,     ["ｱｯﾌﾟﾙストア", "で", "iPhone", "８", "が", "発売", "さ", "れた", "\u3000", "。"]			,     )






							def snake_case_  (   self      ):
									a_     :      List[Any] 						=     MecabTokenizer(normalize_text=a_			,     mecab_dic="ipadic"      )

									self.assertListEqual(
									    tokenizer.tokenize(" \tｱｯﾌﾟﾙストアでiPhone８ が  \n 発売された　。  "      )			,     ["ｱｯﾌﾟﾙストア", "で", "iPhone", "８", "が", "発売", "さ", "れ", "た", "　", "。"]			,     )






							@require_sudachi
							def snake_case_  (   self      ):
									a_     :      int 						=     self.tokenizer_class(self.vocab_file			,     word_tokenizer_type="sudachi"      )
									self.assertIsNotNone(a_      )

									a_     :      Optional[Any] 						=     "こんにちは、世界。\nこんばんは、世界。"
									a_     :      str 						=     tokenizer.tokenize(a_      )
									self.assertListEqual(a_			,     ["こんにちは", "、", "世界", "。", "こん", "##ばんは", "、", "世界", "。"]      )
									self.assertListEqual(tokenizer.convert_tokens_to_ids(a_      )			,     [3, 1_2, 1_0, 1_4, 4, 9, 1_2, 1_0, 1_4]      )

									a_     :      Dict 						=     os.path.join(self.tmpdirname			,     "tokenizer.bin"      )
									with open(a_			,     "wb"      ) as handle:
											pickle.dump(a_			,     a_      )

									with open(a_			,     "rb"      ) as handle:
											a_     :      Optional[Any] 						=     pickle.load(a_      )

									a_     :      int 						=     tokenizer_new.tokenize(a_      )

									self.assertListEqual(a_			,     a_      )






							@require_sudachi
							def snake_case_  (   self      ):
									a_     :      int 						=     SudachiTokenizer(sudachi_dict_type="core"      )

									self.assertListEqual(
									    tokenizer.tokenize(" \tｱｯﾌﾟﾙストアでiPhone８ が  \n 発売された　。  "      )			,     [" ", "\t", "アップル", "ストア", "で", "iPhone", "8", " ", "が", " ", " ", "\n ", "発売", "さ", "れ", "た", " ", "。", " ", " "]			,     )






							@require_sudachi
							def snake_case_  (   self      ):
									a_     :      Any 						=     SudachiTokenizer(sudachi_dict_type="core"			,     sudachi_split_mode="A"      )

									self.assertListEqual(tokenizer.tokenize("外国人参政権"      )			,     ["外国", "人", "参政", "権"]      )






							@require_sudachi
							def snake_case_  (   self      ):
									a_     :      Union[str, Any] 						=     SudachiTokenizer(sudachi_dict_type="core"			,     sudachi_split_mode="B"      )

									self.assertListEqual(tokenizer.tokenize("外国人参政権"      )			,     ["外国人", "参政権"]      )






							@require_sudachi
							def snake_case_  (   self      ):
									a_     :      int 						=     SudachiTokenizer(sudachi_dict_type="core"			,     sudachi_split_mode="C"      )

									self.assertListEqual(tokenizer.tokenize("外国人参政権"      )			,     ["外国人参政権"]      )






							@require_sudachi
							def snake_case_  (   self      ):
									a_     :      List[str] 						=     SudachiTokenizer(do_lower_case=a_			,     sudachi_dict_type="core"      )

									self.assertListEqual(
									    tokenizer.tokenize(" \tｱｯﾌﾟﾙストアでiPhone８ が  \n 発売された　。  "      )			,     [" ", "\t", "アップル", "ストア", "で", "iphone", "8", " ", "が", " ", " ", "\n ", "発売", "さ", "れ", "た", " ", "。", " ", " "]			,     )






							@require_sudachi
							def snake_case_  (   self      ):
									a_     :      str 						=     SudachiTokenizer(normalize_text=a_			,     sudachi_dict_type="core"      )

									self.assertListEqual(
									    tokenizer.tokenize(" \tｱｯﾌﾟﾙストアでiPhone８ が  \n 発売された　。  "      )			,     [" ", "\t", "ｱｯﾌﾟﾙ", "ストア", "で", "iPhone", "８", " ", "が", " ", " ", "\n ", "発売", "さ", "れ", "た", "\u3000", "。", " ", " "]			,     )






							@require_sudachi
							def snake_case_  (   self      ):
									a_     :      Optional[int] 						=     SudachiTokenizer(trim_whitespace=a_			,     sudachi_dict_type="core"      )

									self.assertListEqual(
									    tokenizer.tokenize(" \tｱｯﾌﾟﾙストアでiPhone８ が  \n 発売された　。  "      )			,     ["アップル", "ストア", "で", "iPhone", "8", "が", "発売", "さ", "れ", "た", "。"]			,     )






							@require_jumanpp
							def snake_case_  (   self      ):
									a_     :      Union[str, Any] 						=     self.tokenizer_class(self.vocab_file			,     word_tokenizer_type="jumanpp"      )
									self.assertIsNotNone(a_      )

									a_     :      Optional[int] 						=     "こんにちは、世界。\nこんばんは、世界。"
									a_     :      Any 						=     tokenizer.tokenize(a_      )
									self.assertListEqual(a_			,     ["こんにちは", "、", "世界", "。", "こん", "##ばんは", "、", "世界", "。"]      )
									self.assertListEqual(tokenizer.convert_tokens_to_ids(a_      )			,     [3, 1_2, 1_0, 1_4, 4, 9, 1_2, 1_0, 1_4]      )

									a_     :      int 						=     os.path.join(self.tmpdirname			,     "tokenizer.bin"      )
									with open(a_			,     "wb"      ) as handle:
											pickle.dump(a_			,     a_      )

									with open(a_			,     "rb"      ) as handle:
											a_     :      Tuple 						=     pickle.load(a_      )

									a_     :      Union[str, Any] 						=     tokenizer_new.tokenize(a_      )

									self.assertListEqual(a_			,     a_      )






							@require_jumanpp
							def snake_case_  (   self      ):
									a_     :      Any 						=     JumanppTokenizer()

									self.assertListEqual(
									    tokenizer.tokenize(" \tｱｯﾌﾟﾙストアでiPhone８ が  \n 発売された　。  "      )			,     ["アップル", "ストア", "で", "iPhone", "8", "\u3000", "が", "\u3000", "\u3000", "\u3000", "発売", "さ", "れた", "\u3000", "。"]			,     )






							@require_jumanpp
							def snake_case_  (   self      ):
									a_     :      Optional[Any] 						=     JumanppTokenizer(do_lower_case=a_      )

									self.assertListEqual(
									    tokenizer.tokenize(" \tｱｯﾌﾟﾙストアでiPhone８ が  \n 発売された　。  "      )			,     ["アップル", "ストア", "で", "iphone", "8", "\u3000", "が", "\u3000", "\u3000", "\u3000", "発売", "さ", "れた", "\u3000", "。"]			,     )






							@require_jumanpp
							def snake_case_  (   self      ):
									a_     :      int 						=     JumanppTokenizer(normalize_text=a_      )

									self.assertListEqual(
									    tokenizer.tokenize(" \tｱｯﾌﾟﾙストアでiPhone８ が  \n 発売された　。  "      )			,     ["ｱ", "ｯ", "ﾌ", "ﾟ", "ﾙ", "ストア", "で", "iPhone", "８", "\u3000", "が", "\u3000", "\u3000", "\u3000", "発売", "さ", "れた", "\u3000", "。"]			,     )






							@require_jumanpp
							def snake_case_  (   self      ):
									a_     :      Optional[Any] 						=     JumanppTokenizer(trim_whitespace=a_      )

									self.assertListEqual(
									    tokenizer.tokenize(" \tｱｯﾌﾟﾙストアでiPhone８ が  \n 発売された　。  "      )			,     ["アップル", "ストア", "で", "iPhone", "8", "が", "発売", "さ", "れた", "。"]			,     )






							@require_jumanpp
							def snake_case_  (   self      ):
									a_     :      Any 						=     JumanppTokenizer()

									self.assertListEqual(
									    tokenizer.tokenize("ありがとうございますm(_ _)ｍ見つけるのが大変です。"      )			,     ["ありがとう", "ございます", "m(_ _)m", "見つける", "の", "が", "大変です", "。"]			,     )






							def snake_case_  (   self      ):
									a_     :      str 						=     ["[UNK]", "[CLS]", "[SEP]", "こんにちは", "こん", "にちは", "ばんは", "##こん", "##にちは", "##ばんは"]

									a_     :      int 						=     {}
									for i, token in enumerate(a_      ):
											a_     :      Optional[int] 						=     i
									a_     :      Optional[int] 						=     WordpieceTokenizer(vocab=a_			,     unk_token="[UNK]"      )

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

									self.assertListEqual(tokenizer.tokenize("こんにちは"      )			,     ["こんにちは"]      )

									self.assertListEqual(tokenizer.tokenize("こんばんは"      )			,     ["こん", "##ばんは"]      )

									self.assertListEqual(tokenizer.tokenize("こんばんは こんばんにちは こんにちは"      )			,     ["こん", "##ばんは", "[UNK]", "こんにちは"]      )






							def snake_case_  (   self      ):
									a_     :      Union[str, Any] 						=     BertJapaneseTokenizer.from_pretrained("nlp-waseda/roberta-base-japanese-with-auto-jumanpp"      )
									a_     :      Any 						=     tokenizer.subword_tokenizer

									a_     :      Union[str, Any] 						=     subword_tokenizer.tokenize("国境 の 長い トンネル を 抜ける と 雪国 であった 。"      )
									self.assertListEqual(a_			,     ["▁国境", "▁の", "▁長い", "▁トンネル", "▁を", "▁抜ける", "▁と", "▁雪", "国", "▁であった", "▁。"]      )

									a_     :      int 						=     subword_tokenizer.tokenize("こんばんは こんばん にち は こんにちは"      )
									self.assertListEqual(a_			,     ["▁こん", "ばん", "は", "▁こん", "ばん", "▁に", "ち", "▁は", "▁こんにちは"]      )






							def snake_case_  (   self      ):
									a_     :      List[str] 						=     self.tokenizer_class.from_pretrained("cl-tohoku/bert-base-japanese"      )

									a_     :      Any 						=     tokenizer.encode("ありがとう。"			,     add_special_tokens=a_      )
									a_     :      Optional[Any] 						=     tokenizer.encode("どういたしまして。"			,     add_special_tokens=a_      )

									a_     :      Union[str, Any] 						=     tokenizer.build_inputs_with_special_tokens(a_      )
									a_     :      Tuple 						=     tokenizer.build_inputs_with_special_tokens(a_			,     a_      )

									# 2 is for "[CLS]", 3 is for "[SEP]"
									assert encoded_sentence == [2] + text + [3]
									assert encoded_pair == [2] + text + [3] + text_a + [3]




@custom_tokenizers
class 			snake_case_						(						a_     ,unittest.TestCase				):
							__lowerCAmelCase    					=					BertJapaneseTokenizer
							__lowerCAmelCase    					=					False






							def snake_case_  (   self      ):
									super().setUp()

									a_     :      Any 						=     ["[UNK]", "[CLS]", "[SEP]", "こ", "ん", "に", "ち", "は", "ば", "世", "界", "、", "。"]

									a_     :      str 						=     os.path.join(self.tmpdirname			,     VOCAB_FILES_NAMES["vocab_file"]      )
									with open(self.vocab_file			,     "w"			,     encoding="utf-8"      ) as vocab_writer:
											vocab_writer.write("".join([x + "\n" for x in vocab_tokens]      )      )






							def snake_case_  (   self			,     **a_      ):
									return BertJapaneseTokenizer.from_pretrained(self.tmpdirname			,     subword_tokenizer_type="character"			,     **a_      )






							def snake_case_  (   self			,     a_      ):
									a_     :      int 						=     "こんにちは、世界。 \nこんばんは、世界。"
									a_     :      List[str] 						=     "こ ん に ち は 、 世 界 。 こ ん ば ん は 、 世 界 。"
									return input_text, output_text






							def snake_case_  (   self      ):
									pass  # TODO add if relevant






							def snake_case_  (   self      ):
									pass  # TODO add if relevant






							def snake_case_  (   self      ):
									pass  # TODO add if relevant






							def snake_case_  (   self      ):
									a_     :      Dict 						=     self.tokenizer_class(self.vocab_file			,     subword_tokenizer_type="character"      )

									a_     :      int 						=     tokenizer.tokenize("こんにちは、世界。 \nこんばんは、世界。"      )
									self.assertListEqual(
									    a_			,     ["こ", "ん", "に", "ち", "は", "、", "世", "界", "。", "こ", "ん", "ば", "ん", "は", "、", "世", "界", "。"]      )
									self.assertListEqual(
									    tokenizer.convert_tokens_to_ids(a_      )			,     [3, 4, 5, 6, 7, 1_1, 9, 1_0, 1_2, 3, 4, 8, 4, 7, 1_1, 9, 1_0, 1_2]      )






							def snake_case_  (   self      ):
									a_     :      Union[str, Any] 						=     ["[UNK]", "[CLS]", "[SEP]", "こ", "ん", "に", "ち", "は", "ば", "世", "界", "、", "。"]

									a_     :      Any 						=     {}
									for i, token in enumerate(a_      ):
											a_     :      str 						=     i
									a_     :      List[str] 						=     CharacterTokenizer(vocab=a_			,     unk_token="[UNK]"      )

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

									self.assertListEqual(tokenizer.tokenize("こんにちは"      )			,     ["こ", "ん", "に", "ち", "は"]      )

									self.assertListEqual(tokenizer.tokenize("こんにちほ"      )			,     ["こ", "ん", "に", "ち", "[UNK]"]      )






							def snake_case_  (   self      ):
									a_     :      str 						=     self.tokenizer_class.from_pretrained("cl-tohoku/bert-base-japanese-char"      )

									a_     :      Optional[int] 						=     tokenizer.encode("ありがとう。"			,     add_special_tokens=a_      )
									a_     :      Tuple 						=     tokenizer.encode("どういたしまして。"			,     add_special_tokens=a_      )

									a_     :      Optional[Any] 						=     tokenizer.build_inputs_with_special_tokens(a_      )
									a_     :      Any 						=     tokenizer.build_inputs_with_special_tokens(a_			,     a_      )

									# 2 is for "[CLS]", 3 is for "[SEP]"
									assert encoded_sentence == [2] + text + [3]
									assert encoded_pair == [2] + text + [3] + text_a + [3]




@custom_tokenizers
class 			snake_case_						(						unittest.TestCase				):






							def snake_case_  (   self      ):
									a_     :      List[Any] 						=     "cl-tohoku/bert-base-japanese"
									a_     :      Dict 						=     AutoTokenizer.from_pretrained(a_      )
									self.assertIsInstance(a_			,     a_      )




class 			snake_case_						(						unittest.TestCase				):






							def snake_case_  (   self      ):
									a_     :      Any 						=     "cl-tohoku/bert-base-japanese"
									with self.assertLogs("transformers"			,     level="WARNING"      ) as cm:
											BertTokenizer.from_pretrained(a_      )
											self.assertTrue(
											    cm.records[0].message.startswith(
											        "The tokenizer class you load from this checkpoint is not the same type as the class this function"
											        " is called from."      )      )
									a_     :      Any 						=     "bert-base-cased"
									with self.assertLogs("transformers"			,     level="WARNING"      ) as cm:
											BertJapaneseTokenizer.from_pretrained(a_      )
											self.assertTrue(
											    cm.records[0].message.startswith(
											        "The tokenizer class you load from this checkpoint is not the same type as the class this function"
											        " is called from."      )      )




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




@flax.struct.dataclass
class 			snake_case_						(						a_				):
							__lowerCAmelCase    					=					42
							__lowerCAmelCase    					=					42




class 			snake_case_						(						nn.Module				):
							__lowerCAmelCase    					=					42
							__lowerCAmelCase    					=					(1_6, 3_2, 9_6, 2_5_6)
							__lowerCAmelCase    					=					jnp.floataa






							def snake_case_  (   self      ):
									a_     :      Any 						=     nn.Conv(
									    self.block_out_channels[0]			,     kernel_size=(3, 3)			,     padding=((1, 1), (1, 1))			,     dtype=self.dtype			,     )

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

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

									for block in self.blocks:
											a_     :      Tuple 						=     block(a_      )
											a_     :      Union[str, Any] 						=     nn.silu(a_      )

									a_     :      List[str] 						=     self.conv_out(a_      )

									return embedding




@flax_register_to_config
class 			snake_case_						(						nn.Module     ,a_     ,a_				):
							__lowerCAmelCase    					=					3_2
							__lowerCAmelCase    					=					4
							__lowerCAmelCase    					=					(
							    "CrossAttnDownBlock2D",
							    "CrossAttnDownBlock2D",
							    "CrossAttnDownBlock2D",
							    "DownBlock2D",
							)
							__lowerCAmelCase    					=					False
							__lowerCAmelCase    					=					(3_2_0, 6_4_0, 1_2_8_0, 1_2_8_0)
							__lowerCAmelCase    					=					2
							__lowerCAmelCase    					=					8
							__lowerCAmelCase    					=					None
							__lowerCAmelCase    					=					1_2_8_0
							__lowerCAmelCase    					=					0.0
							__lowerCAmelCase    					=					False
							__lowerCAmelCase    					=					jnp.floataa
							__lowerCAmelCase    					=					True
							__lowerCAmelCase    					=					0
							__lowerCAmelCase    					=					"rgb"
							__lowerCAmelCase    					=					(1_6, 3_2, 9_6, 2_5_6)






							def snake_case_  (   self			,     a_      ):
									# init input tensors
									a_     :      Any 						=     (1, self.in_channels, self.sample_size, self.sample_size)
									a_     :      Tuple 						=     jnp.zeros(a_			,     dtype=jnp.floataa      )
									a_     :      Optional[int] 						=     jnp.ones((1,)			,     dtype=jnp.intaa      )
									a_     :      Any 						=     jnp.zeros((1, 1, self.cross_attention_dim)			,     dtype=jnp.floataa      )
									a_     :      Union[str, Any] 						=     (1, 3, self.sample_size * 8, self.sample_size * 8)
									a_     :      Optional[int] 						=     jnp.zeros(a_			,     dtype=jnp.floataa      )

									a_						,				a_     :      List[Any] 						=     jax.random.split(a_      )
									a_     :      Dict 						=     {"params": params_rng, "dropout": dropout_rng}

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






							def snake_case_  (   self      ):
									a_     :      List[Any] 						=     self.block_out_channels
									a_     :      Any 						=     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.
									a_     :      Optional[int] 						=     self.num_attention_heads or self.attention_head_dim

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

									# time
									a_     :      Optional[Any] 						=     FlaxTimesteps(
									    block_out_channels[0]			,     flip_sin_to_cos=self.flip_sin_to_cos			,     freq_shift=self.config.freq_shift      )
									a_     :      List[str] 						=     FlaxTimestepEmbedding(a_			,     dtype=self.dtype      )

									a_     :      Optional[int] 						=     FlaxControlNetConditioningEmbedding(
									    conditioning_embedding_channels=block_out_channels[0]			,     block_out_channels=self.conditioning_embedding_out_channels			,     )

									a_     :      int 						=     self.only_cross_attention
									if isinstance(a_			,     a_      ):
											a_     :      Tuple 						=     (only_cross_attention,) * len(self.down_block_types      )

									if isinstance(a_			,     a_      ):
											a_     :      int 						=     (num_attention_heads,) * len(self.down_block_types      )

									# down
									a_     :      Any 						=     []
									a_     :      Optional[Any] 						=     []

									a_     :      List[Any] 						=     block_out_channels[0]

									a_     :      Dict 						=     nn.Conv(
									    a_			,     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(a_      )

									for i, down_block_type in enumerate(self.down_block_types      ):
											a_     :      int 						=     output_channel
											a_     :      Optional[int] 						=     block_out_channels[i]
											a_     :      str 						=     i == len(a_      ) - 1

											if down_block_type == "CrossAttnDownBlock2D":
													a_     :      int 						=     FlaxCrossAttnDownBlockaD(
													    in_channels=a_			,     out_channels=a_			,     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:
													a_     :      Union[str, Any] 						=     FlaxDownBlockaD(
													    in_channels=a_			,     out_channels=a_			,     dropout=self.dropout			,     num_layers=self.layers_per_block			,     add_downsample=not is_final_block			,     dtype=self.dtype			,     )

											down_blocks.append(a_      )

											for _ in range(self.layers_per_block      ):
													a_     :      Tuple 						=     nn.Conv(
													    a_			,     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(a_      )

											if not is_final_block:
													a_     :      Optional[Any] 						=     nn.Conv(
													    a_			,     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(a_      )

									a_     :      Union[str, Any] 						=     down_blocks
									a_     :      Optional[int] 						=     controlnet_down_blocks

									# mid
									a_     :      Union[str, Any] 						=     block_out_channels[-1]
									a_     :      Any 						=     FlaxUNetMidBlockaDCrossAttn(
									    in_channels=a_			,     dropout=self.dropout			,     num_attention_heads=num_attention_heads[-1]			,     use_linear_projection=self.use_linear_projection			,     dtype=self.dtype			,     )

									a_     :      Union[str, Any] 						=     nn.Conv(
									    a_			,     kernel_size=(1, 1)			,     padding="VALID"			,     kernel_init=nn.initializers.zeros_init()			,     bias_init=nn.initializers.zeros_init()			,     dtype=self.dtype			,     )






							def __call__(   self			,     a_			,     a_			,     a_			,     a_			,     a_ = 1.0			,     a_ = True			,     a_ = False			,     ):
									a_     :      int 						=     self.controlnet_conditioning_channel_order
									if channel_order == "bgr":
											a_     :      Tuple 						=     jnp.flip(a_			,     axis=1      )

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

									a_     :      List[str] 						=     self.time_proj(a_      )
									a_     :      Union[str, Any] 						=     self.time_embedding(a_      )

									# 2. pre-process
									a_     :      Union[str, Any] 						=     jnp.transpose(a_			,     (0, 2, 3, 1)      )
									a_     :      Union[str, Any] 						=     self.conv_in(a_      )

									a_     :      Optional[int] 						=     jnp.transpose(a_			,     (0, 2, 3, 1)      )
									a_     :      int 						=     self.controlnet_cond_embedding(a_      )
									sample += controlnet_cond

									# 3. down
									a_     :      List[str] 						=     (sample,)
									for down_block in self.down_blocks:
											if isinstance(a_			,     a_      ):
													a_						,				a_     :      int 						=     down_block(a_			,     a_			,     a_			,     deterministic=not train      )
											else:
													a_						,				a_     :      Union[str, Any] 						=     down_block(a_			,     a_			,     deterministic=not train      )
											down_block_res_samples += res_samples

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

									# 5. contronet blocks
									a_     :      Optional[int] 						=     ()
									for down_block_res_sample, controlnet_block in zip(a_			,     self.controlnet_down_blocks      ):
											a_     :      Tuple 						=     controlnet_block(a_      )
											controlnet_down_block_res_samples += (down_block_res_sample,)

									a_     :      Union[str, Any] 						=     controlnet_down_block_res_samples

									a_     :      Optional[int] 						=     self.controlnet_mid_block(a_      )

									# 6. scaling
									a_     :      Optional[Any] 						=     [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=a_			,     mid_block_res_sample=a_      )