




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 					lowerCAmelCase_ (       _a			,unittest.TestCase):
						lowerCamelCase_         =     BertJapaneseTokenizer
						lowerCamelCase_         =     False
						lowerCamelCase_         =     True




						def  _snake_case					(	self						:   Optional[int]						)      ->str:






									"""simple docstring"""






									super().setUp()

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

									a__ :Any						  =    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						:   List[Any]   ,     __A						:   Optional[int]						)      ->List[Any]:






									"""simple docstring"""






									a__ :Any						  =    "こんにちは、世界。 \nこんばんは、世界。"
									a__ :List[str]						  =    "こんにちは 、 世界 。 こんばんは 、 世界 。"
									return input_text, output_text




						def  _snake_case					(	self						:   int   ,     __A						:   Any						)      ->List[Any]:






									"""simple docstring"""






									a__ :List[str]						  =    self.get_input_output_texts(__A						)
									a__ :List[Any]						  =    tokenizer.encode(__A   ,     add_special_tokens=__A						)
									a__ :Union[str, Any]						  =    tokenizer.decode(__A   ,     clean_up_tokenization_spaces=__A						)
									return text, ids




						def  _snake_case					(	self						:   List[str]						)      ->Optional[Any]:






									"""simple docstring"""






									pass  # TODO add if relevant




						def  _snake_case					(	self						:   List[str]						)      ->str:






									"""simple docstring"""






									pass  # TODO add if relevant




						def  _snake_case					(	self						:   Dict						)      ->int:






									"""simple docstring"""






									pass  # TODO add if relevant




						def  _snake_case					(	self						:   Tuple						)      ->List[str]:






									"""simple docstring"""






									a__ :Union[str, Any]						  =    self.tokenizer_class(self.vocab_file						)

									a__ :Optional[int]						  =    tokenizer.tokenize("こんにちは、世界。\nこんばんは、世界。"						)
									self.assertListEqual(__A   ,     ["こんにちは", "、", "世界", "。", "こん", "##ばんは", "、", "世界", "。"]						)
									self.assertListEqual(tokenizer.convert_tokens_to_ids(__A						)   ,     [3, 12, 10, 14, 4, 9, 12, 10, 14]						)




						def  _snake_case					(	self						:   List[str]						)      ->Union[str, Any]:






									"""simple docstring"""






									a__ :Dict						  =    self.tokenizer_class(self.vocab_file   ,     word_tokenizer_type="mecab"						)
									self.assertIsNotNone(__A						)

									a__ :Union[str, Any]						  =    "こんにちは、世界。\nこんばんは、世界。"
									a__ :str						  =    tokenizer.tokenize(__A						)
									self.assertListEqual(__A   ,     ["こんにちは", "、", "世界", "。", "こん", "##ばんは", "、", "世界", "。"]						)
									self.assertListEqual(tokenizer.convert_tokens_to_ids(__A						)   ,     [3, 12, 10, 14, 4, 9, 12, 10, 14]						)

									a__ :List[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__ :Optional[int]						  =    tokenizer_new.tokenize(__A						)

									self.assertListEqual(__A   ,     __A						)




						def  _snake_case					(	self						:   List[Any]						)      ->int:






									"""simple docstring"""






									a__ :Any						  =    MecabTokenizer(mecab_dic="ipadic"						)

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




						def  _snake_case					(	self						:   Any						)      ->Optional[int]:






									"""simple docstring"""






									try:
												a__ :List[Any]						  =    MecabTokenizer(mecab_dic="unidic_lite"						)
									except ModuleNotFoundError:
												return

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




						def  _snake_case					(	self						:   Tuple						)      ->Dict:






									"""simple docstring"""






									try:
												a__ :Union[str, Any]						  =    MecabTokenizer(mecab_dic="unidic"						)
									except ModuleNotFoundError:
												return

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




						def  _snake_case					(	self						:   int						)      ->Tuple:






									"""simple docstring"""






									a__ :str						  =    MecabTokenizer(do_lower_case=__A   ,     mecab_dic="ipadic"						)

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




						def  _snake_case					(	self						:   List[str]						)      ->List[Any]:






									"""simple docstring"""






									try:
												a__ :Any						  =    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						:   Optional[int]						)      ->int:






									"""simple docstring"""






									a__ :int						  =    MecabTokenizer(normalize_text=__A   ,     mecab_dic="ipadic"						)

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




						@require_sudachi
						def  _snake_case					(	self						:   Dict						)      ->str:






									"""simple docstring"""






									a__ :int						  =    self.tokenizer_class(self.vocab_file   ,     word_tokenizer_type="sudachi"						)
									self.assertIsNotNone(__A						)

									a__ :List[Any]						  =    "こんにちは、世界。\nこんばんは、世界。"
									a__ :List[str]						  =    tokenizer.tokenize(__A						)
									self.assertListEqual(__A   ,     ["こんにちは", "、", "世界", "。", "こん", "##ばんは", "、", "世界", "。"]						)
									self.assertListEqual(tokenizer.convert_tokens_to_ids(__A						)   ,     [3, 12, 10, 14, 4, 9, 12, 10, 14]						)

									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__ :Tuple						  =    pickle.load(__A						)

									a__ :Any						  =    tokenizer_new.tokenize(__A						)

									self.assertListEqual(__A   ,     __A						)




						@require_sudachi
						def  _snake_case					(	self						:   Union[str, Any]						)      ->Any:






									"""simple docstring"""






									a__ :Dict						  =    SudachiTokenizer(sudachi_dict_type="core"						)

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




						@require_sudachi
						def  _snake_case					(	self						:   Optional[int]						)      ->Tuple:






									"""simple docstring"""






									a__ :Union[str, Any]						  =    SudachiTokenizer(sudachi_dict_type="core"   ,     sudachi_split_mode="A"						)

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




						@require_sudachi
						def  _snake_case					(	self						:   Union[str, Any]						)      ->Tuple:






									"""simple docstring"""






									a__ :Dict						  =    SudachiTokenizer(sudachi_dict_type="core"   ,     sudachi_split_mode="B"						)

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




						@require_sudachi
						def  _snake_case					(	self						:   Any						)      ->Optional[int]:






									"""simple docstring"""






									a__ :List[str]						  =    SudachiTokenizer(sudachi_dict_type="core"   ,     sudachi_split_mode="C"						)

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




						@require_sudachi
						def  _snake_case					(	self						:   Tuple						)      ->Optional[int]:






									"""simple docstring"""






									a__ :Tuple						  =    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						:   int						)      ->Optional[int]:






									"""simple docstring"""






									a__ :Tuple						  =    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						:   Union[str, Any]						)      ->List[Any]:






									"""simple docstring"""






									a__ :int						  =    SudachiTokenizer(trim_whitespace=__A   ,     sudachi_dict_type="core"						)

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




						@require_jumanpp
						def  _snake_case					(	self						:   Tuple						)      ->List[str]:






									"""simple docstring"""






									a__ :Tuple						  =    self.tokenizer_class(self.vocab_file   ,     word_tokenizer_type="jumanpp"						)
									self.assertIsNotNone(__A						)

									a__ :Optional[int]						  =    "こんにちは、世界。\nこんばんは、世界。"
									a__ :str						  =    tokenizer.tokenize(__A						)
									self.assertListEqual(__A   ,     ["こんにちは", "、", "世界", "。", "こん", "##ばんは", "、", "世界", "。"]						)
									self.assertListEqual(tokenizer.convert_tokens_to_ids(__A						)   ,     [3, 12, 10, 14, 4, 9, 12, 10, 14]						)

									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__ :Optional[Any]						  =    pickle.load(__A						)

									a__ :str						  =    tokenizer_new.tokenize(__A						)

									self.assertListEqual(__A   ,     __A						)




						@require_jumanpp
						def  _snake_case					(	self						:   Optional[int]						)      ->List[str]:






									"""simple docstring"""






									a__ :Tuple						  =    JumanppTokenizer()

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




						@require_jumanpp
						def  _snake_case					(	self						:   Any						)      ->Optional[int]:






									"""simple docstring"""






									a__ :Union[str, 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						:   str						)      ->List[str]:






									"""simple docstring"""






									a__ :Dict						  =    JumanppTokenizer(normalize_text=__A						)

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




						@require_jumanpp
						def  _snake_case					(	self						:   Optional[int]						)      ->int:






									"""simple docstring"""






									a__ :Optional[int]						  =    JumanppTokenizer(trim_whitespace=__A						)

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




						@require_jumanpp
						def  _snake_case					(	self						:   Optional[int]						)      ->Any:






									"""simple docstring"""






									a__ :List[str]						  =    JumanppTokenizer()

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




						def  _snake_case					(	self						:   Tuple						)      ->List[Any]:






									"""simple docstring"""






									a__ :str						  =    ["[UNK]", "[CLS]", "[SEP]", "こんにちは", "こん", "にちは", "ばんは", "##こん", "##にちは", "##ばんは"]

									a__ :str						  =    {}
									for i, token in enumerate(__A						):
												a__ :Tuple						  =    i
									a__ :str						  =    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						:   Dict						)      ->Any:






									"""simple docstring"""






									a__ :Optional[Any]						  =    BertJapaneseTokenizer.from_pretrained("nlp-waseda/roberta-base-japanese-with-auto-jumanpp"						)
									a__ :Dict						  =    tokenizer.subword_tokenizer

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

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






						def  _snake_case					(	self						:   Optional[int]						)      ->List[str]:






									"""simple docstring"""






									a__ :Tuple						  =    self.tokenizer_class.from_pretrained("cl-tohoku/bert-base-japanese"						)

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

									a__ :Optional[Any]						  =    tokenizer.build_inputs_with_special_tokens(__A						)
									a__ :List[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 					lowerCAmelCase_ (       _a			,unittest.TestCase):
						lowerCamelCase_         =     BertJapaneseTokenizer
						lowerCamelCase_         =     False




						def  _snake_case					(	self						:   Tuple						)      ->str:






									"""simple docstring"""






									super().setUp()

									a__ :Optional[int]						  =    ["[UNK]", "[CLS]", "[SEP]", "こ", "ん", "に", "ち", "は", "ば", "世", "界", "、", "。"]

									a__ :List[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						:   Optional[Any]   ,     **__A						:   Any						)      ->List[str]:






									"""simple docstring"""






									return BertJapaneseTokenizer.from_pretrained(self.tmpdirname   ,     subword_tokenizer_type="character"   ,     **__A						)




						def  _snake_case					(	self						:   Optional[Any]   ,     __A						:   Optional[int]						)      ->Tuple:






									"""simple docstring"""






									a__ :List[Any]						  =    "こんにちは、世界。 \nこんばんは、世界。"
									a__ :Union[str, Any]						  =    "こ ん に ち は 、 世 界 。 こ ん ば ん は 、 世 界 。"
									return input_text, output_text




						def  _snake_case					(	self						:   Optional[Any]						)      ->Any:






									"""simple docstring"""






									pass  # TODO add if relevant




						def  _snake_case					(	self						:   Optional[int]						)      ->List[str]:






									"""simple docstring"""






									pass  # TODO add if relevant




						def  _snake_case					(	self						:   Dict						)      ->List[str]:






									"""simple docstring"""






									pass  # TODO add if relevant




						def  _snake_case					(	self						:   Union[str, Any]						)      ->int:






									"""simple docstring"""






									a__ :Dict						  =    self.tokenizer_class(self.vocab_file   ,     subword_tokenizer_type="character"						)

									a__ :Dict						  =    tokenizer.tokenize("こんにちは、世界。 \nこんばんは、世界。"						)
									self.assertListEqual(
									    __A   ,     ["こ", "ん", "に", "ち", "は", "、", "世", "界", "。", "こ", "ん", "ば", "ん", "は", "、", "世", "界", "。"]						)
									self.assertListEqual(
									    tokenizer.convert_tokens_to_ids(__A						)   ,     [3, 4, 5, 6, 7, 11, 9, 10, 12, 3, 4, 8, 4, 7, 11, 9, 10, 12]						)




						def  _snake_case					(	self						:   Optional[int]						)      ->int:






									"""simple docstring"""






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

									a__ :Any						  =    {}
									for i, token in enumerate(__A						):
												a__ :Any						  =    i
									a__ :int						  =    CharacterTokenizer(vocab=__A   ,     unk_token="[UNK]"						)

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

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

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






						def  _snake_case					(	self						:   Any						)      ->List[str]:






									"""simple docstring"""






									a__ :Dict						  =    self.tokenizer_class.from_pretrained("cl-tohoku/bert-base-japanese-char"						)

									a__ :int						  =    tokenizer.encode("ありがとう。"   ,     add_special_tokens=__A						)
									a__ :Optional[int]						  =    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 					lowerCAmelCase_ (       unittest.TestCase):




						def  _snake_case					(	self						:   str						)      ->Union[str, Any]:






									"""simple docstring"""






									a__ :Any						  =    "cl-tohoku/bert-base-japanese"
									a__ :Dict						  =    AutoTokenizer.from_pretrained(__A						)
									self.assertIsInstance(__A   ,     __A						)






class 					lowerCAmelCase_ (       unittest.TestCase):




						def  _snake_case					(	self						:   List[Any]						)      ->int:






									"""simple docstring"""






									a__ :List[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."						)						)






import os
from glob import glob

import imageio
import torch
import torchvision
import wandb
from img_processing import custom_to_pil, loop_post_process, preprocess, preprocess_vqgan
from loaders import load_vqgan
from PIL import Image
from torch import nn

from transformers import CLIPModel, CLIPTokenizerFast
from utils import get_device, get_timestamp, show_pil


class 					lowerCAmelCase_ :




						def __init__(	self						:   str   ,     __A						:   str = "cpu"   ,     __A						:   str = "openai/clip-vit-large-patch14"						)      ->None:






									"""simple docstring"""






									a__ :List[str]						  =    device
									a__ :Optional[Any]						  =    CLIPTokenizerFast.from_pretrained(__A						)
									a__ :str						  =    [0.48_145_466, 0.4_578_275, 0.40_821_073]
									a__ :Any						  =    [0.26_862_954, 0.26_130_258, 0.27_577_711]
									a__ :Union[str, Any]						  =    torchvision.transforms.Normalize(self.image_mean   ,     self.image_std						)
									a__ :List[Any]						  =    torchvision.transforms.Resize(224						)
									a__ :Optional[Any]						  =    torchvision.transforms.CenterCrop(224						)




						def  _snake_case					(	self						:   int   ,     __A						:   int						)      ->List[Any]:






									"""simple docstring"""






									a__ :List[str]						  =    self.resize(__A						)
									a__ :Union[str, Any]						  =    self.center_crop(__A						)
									a__ :List[str]						  =    self.normalize(__A						)
									return images






						def __call__(	self						:   Union[str, Any]   ,     __A						:   Dict=None   ,     __A						:   List[Any]=None   ,     **__A						:   Union[str, Any]						)      ->Tuple:






									"""simple docstring"""






									a__ :Optional[int]						  =    self.tokenizer(text=__A   ,     **__A						)
									a__ :str						  =    self.preprocess_img(__A						)
									a__ :Any						  =    {key: value.to(self.device						) for (key, value) in encoding.items()}
									return encoding






class 					lowerCAmelCase_ (       nn.Module):




						def __init__(	self						:   List[str]   ,     __A						:   List[str]=10   ,     __A						:   List[Any]=0.01   ,     __A						:   Optional[int]=None   ,     __A						:   List[str]=None   ,     __A						:   Dict=None   ,     __A						:   List[str]=None   ,     __A						:   Tuple=None   ,     __A						:   str=None   ,     __A						:   Any=False   ,     __A						:   Union[str, Any]=True   ,     __A						:   Tuple="image"   ,     __A						:   Tuple=True   ,     __A						:   int=False   ,     __A						:   Tuple=False   ,     __A						:   Union[str, Any]=False   ,     )      ->None:






									"""simple docstring"""






									super().__init__()
									a__ :List[str]						  =    None
									a__ :int						  =    device if device else get_device()
									if vqgan:
												a__ :Any						  =    vqgan
									else:
												a__ :Optional[int]						  =    load_vqgan(self.device   ,     conf_path=__A   ,     ckpt_path=__A						)
									self.vqgan.eval()
									if clip:
												a__ :Dict						  =    clip
									else:
												a__ :Any						  =    CLIPModel.from_pretrained("openai/clip-vit-base-patch32"						)
									self.clip.to(self.device						)
									a__ :Optional[int]						  =    ProcessorGradientFlow(device=self.device						)

									a__ :Any						  =    iterations
									a__ :Any						  =    lr
									a__ :List[str]						  =    log
									a__ :List[Any]						  =    make_grid
									a__ :Optional[int]						  =    return_val
									a__ :Optional[Any]						  =    quantize
									a__ :Any						  =    self.vqgan.decoder.z_shape




						def  _snake_case					(	self						:   str   ,     __A						:   str=None   ,     __A						:   Optional[int]=None   ,     __A						:   int=5   ,     __A						:   Any=True						)      ->Optional[int]:






									"""simple docstring"""






									a__ :Optional[int]						  =    []
									if output_path is None:
												a__ :Any						  =    "./animation.gif"
									if input_path is None:
												a__ :Dict						  =    self.save_path
									a__ :Dict						  =    sorted(glob(input_path + "/*"						)						)
									if not len(__A						):
												raise ValueError(
												    "No images found in save path, aborting (did you pass save_intermediate=True to the generate"
												    " function?)"						)
									if len(__A						) == 1:
												print("Only one image found in save path, (did you pass save_intermediate=True to the generate function?)"						)
									a__ :str						  =    total_duration / len(__A						)
									a__ :str						  =    [frame_duration] * len(__A						)
									if extend_frames:
												a__ :Optional[Any]						  =    1.5
												a__ :Optional[Any]						  =    3
									for file_name in paths:
												if file_name.endswith(".png"						):
															images.append(imageio.imread(__A						)						)
									imageio.mimsave(__A   ,     __A   ,     duration=__A						)
									print(F'''gif saved to {output_path}'''						)




						def  _snake_case					(	self						:   str   ,     __A						:   Tuple=None   ,     __A						:   Optional[int]=None						)      ->Optional[int]:






									"""simple docstring"""






									if not (path or img):
												raise ValueError("Input either path or tensor"						)
									if img is not None:
												raise NotImplementedError
									a__ :List[Any]						  =    preprocess(Image.open(__A						)   ,     target_image_size=256						).to(self.device						)
									a__ :List[str]						  =    preprocess_vqgan(__A						)
									a__  ,  *a__ :Union[str, Any]						  =    self.vqgan.encode(__A						)
									return z




						def  _snake_case					(	self						:   List[str]   ,     __A						:   Any						)      ->Optional[int]:






									"""simple docstring"""






									a__ :Dict						  =    self.latent.detach().requires_grad_()
									a__ :int						  =    base_latent + transform_vector
									if self.quantize:
												a__  ,  *a__ :str						  =    self.vqgan.quantize(__A						)
									else:
												a__ :List[str]						  =    trans_latent
									return self.vqgan.decode(__A						)




						def  _snake_case					(	self						:   List[str]   ,     __A						:   Optional[Any]   ,     __A						:   Optional[Any]   ,     __A						:   List[Any]=None						)      ->str:






									"""simple docstring"""






									a__ :Dict						  =    self.clip_preprocessor(text=__A   ,     images=__A   ,     return_tensors="pt"   ,     padding=__A						)
									a__ :List[str]						  =    self.clip(**__A						)
									a__ :Optional[int]						  =    clip_outputs.logits_per_image
									if weights is not None:
												a__ :Tuple						  =    similarity_logits * weights
									return similarity_logits.sum()




						def  _snake_case					(	self						:   List[str]   ,     __A						:   List[Any]   ,     __A						:   Optional[int]   ,     __A						:   Optional[int]						)      ->int:






									"""simple docstring"""






									a__ :str						  =    self._get_clip_similarity(pos_prompts["prompts"]   ,     __A   ,     weights=(1 / pos_prompts["weights"])						)
									if neg_prompts:
												a__ :Optional[int]						  =    self._get_clip_similarity(neg_prompts["prompts"]   ,     __A   ,     weights=neg_prompts["weights"]						)
									else:
												a__ :Tuple						  =    torch.tensor([1]   ,     device=self.device						)
									a__ :Dict						  =    -torch.log(__A						) + torch.log(__A						)
									return loss




						def  _snake_case					(	self						:   int   ,     __A						:   Dict   ,     __A						:   Union[str, Any]   ,     __A						:   List[str]						)      ->List[str]:






									"""simple docstring"""






									a__ :List[Any]						  =    torch.randn_like(self.latent   ,     requires_grad=__A   ,     device=self.device						)
									a__ :Optional[Any]						  =    torch.optim.Adam([vector]   ,     lr=self.lr						)

									for i in range(self.iterations						):
												optim.zero_grad()
												a__ :Tuple						  =    self._add_vector(__A						)
												a__ :Dict						  =    loop_post_process(__A						)
												a__ :str						  =    self._get_CLIP_loss(__A   ,     __A   ,     __A						)
												print("CLIP loss"   ,     __A						)
												if self.log:
															wandb.log({"CLIP Loss": clip_loss}						)
												clip_loss.backward(retain_graph=__A						)
												optim.step()
												if self.return_val == "image":
															yield custom_to_pil(transformed_img[0]						)
												else:
															yield vector




						def  _snake_case					(	self						:   Union[str, Any]   ,     __A						:   Optional[Any]   ,     __A						:   Union[str, Any]   ,     __A						:   Optional[Any]						)      ->Optional[int]:






									"""simple docstring"""






									wandb.init(reinit=__A   ,     project="face-editor"						)
									wandb.config.update({"Positive Prompts": positive_prompts}						)
									wandb.config.update({"Negative Prompts": negative_prompts}						)
									wandb.config.update({"lr": self.lr, "iterations": self.iterations}						)
									if image_path:
												a__ :List[str]						  =    Image.open(__A						)
												a__ :Any						  =    image.resize((256, 256)						)
												wandb.log("Original Image"   ,     wandb.Image(__A						)						)




						def  _snake_case					(	self						:   Union[str, Any]   ,     __A						:   Union[str, Any]						)      ->int:






									"""simple docstring"""






									if not prompts:
												return []
									a__ :List[str]						  =    []
									a__ :Tuple						  =    []
									if isinstance(__A   ,     __A						):
												a__ :Tuple						  =    [prompt.strip() for prompt in prompts.split("|"						)]
									for prompt in prompts:
												if isinstance(__A   ,     (tuple, list)						):
															a__ :str						  =    prompt[0]
															a__ :List[Any]						  =    float(prompt[1]						)
												elif ":" in prompt:
															a__  ,  a__ :Tuple						  =    prompt.split(":"						)
															a__ :Dict						  =    float(__A						)
												else:
															a__ :int						  =    prompt
															a__ :Any						  =    1.0
												processed_prompts.append(__A						)
												weights.append(__A						)
									return {
									    "prompts": processed_prompts,
									    "weights": torch.tensor(__A   ,     device=self.device						),
									}






						def  _snake_case					(	self						:   Optional[Any]   ,     __A						:   Any   ,     __A						:   int=None   ,     __A						:   int=None   ,     __A						:   str=True   ,     __A						:   List[Any]=False   ,     __A						:   Optional[int]=True   ,     __A						:   Any=True   ,     __A						:   Optional[Any]=None   ,     )      ->Optional[Any]:






									"""simple docstring"""






									if image_path:
												a__ :Optional[int]						  =    self._get_latent(__A						)
									else:
												a__ :Dict						  =    torch.randn(self.latent_dim   ,     device=self.device						)
									if self.log:
												self._init_logging(__A   ,     __A   ,     __A						)

									assert pos_prompts, "You must provide at least one positive prompt."
									a__ :int						  =    self.process_prompts(__A						)
									a__ :Any						  =    self.process_prompts(__A						)
									if save_final and save_path is None:
												a__ :Dict						  =    os.path.join("./outputs/"   ,     "_".join(pos_prompts["prompts"]						)						)
									if not os.path.exists(__A						):
												os.makedirs(__A						)
									else:
												a__ :str						  =    save_path + "_" + get_timestamp()
												os.makedirs(__A						)
									a__ :Any						  =    save_path

									a__ :Optional[int]						  =    self.vqgan.decode(self.latent						)[0]
									if show_intermediate:
												print("Original Image"						)
												show_pil(custom_to_pil(__A						)						)

									a__ :List[str]						  =    loop_post_process(__A						)
									for iter, transformed_img in enumerate(self._optimize_CLIP(__A   ,     __A   ,     __A						)						):
												if show_intermediate:
															show_pil(__A						)
												if save_intermediate:
															transformed_img.save(os.path.join(self.save_path   ,     F'''iter_{iter:03d}.png'''						)						)
												if self.log:
															wandb.log({"Image": wandb.Image(__A						)}						)
									if show_final:
												show_pil(__A						)
									if save_final:
												transformed_img.save(os.path.join(self.save_path   ,     F'''iter_{iter:03d}_final.png'''						)						)


