

import os
import tempfile
import unittest

import numpy as np
import torch

from diffusers.models import ModelMixin, UNetaDConditionModel
from diffusers.models.attention_processor import AttnProcessor, LoRAAttnProcessor
from diffusers.utils import (
    floats_tensor,
    logging,
    skip_mps,
    torch_device,
)
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin, UNetTesterMixin


enable_full_determinism()

lowerCAmelCase_	     =      logging.get_logger(__name__)


def   __SCREAMING_SNAKE_CASE (SCREAMING_SNAKE_CASE__  ,		SCREAMING_SNAKE_CASE__ = True	):
						snake_case_       =     {}
						for name in model.attn_processors.keys():
												snake_case_       =     name.endswith('''attn2.processor'''	) and not (
												    name.startswith('''transformer_in'''	) or "temp_attentions" in name.split('''.'''	)
												)
												snake_case_       =     model.config.cross_attention_dim if has_cross_attention else None
												if name.startswith('''mid_block'''	):
																		snake_case_       =     model.config.block_out_channels[-1]
												elif name.startswith('''up_blocks'''	):
																		snake_case_       =     int(name[len('''up_blocks.'''	)]	)
																		snake_case_       =     list(reversed(model.config.block_out_channels	)	)[block_id]
												elif name.startswith('''down_blocks'''	):
																		snake_case_       =     int(name[len('''down_blocks.'''	)]	)
																		snake_case_       =     model.config.block_out_channels[block_id]
												elif name.startswith('''transformer_in'''	):
																		# Note that the `8 * ...` comes from: https://github.com/huggingface/diffusers/blob/7139f0e874f10b2463caa8cbd585762a309d12d6/src/diffusers/models/unet_3d_condition.py#L148
																		snake_case_       =     8 * model.config.attention_head_dim

												snake_case_       =     LoRAAttnProcessor(hidden_size=_SCREAMING_SNAKE_CASE  ,		cross_attention_dim=_SCREAMING_SNAKE_CASE	)
												snake_case_       =     lora_attn_procs[name].to(model.device	)

												if mock_weights:
																		# add 1 to weights to mock trained weights
																		with torch.no_grad():
																								lora_attn_procs[name].to_q_lora.up.weight += 1
																								lora_attn_procs[name].to_k_lora.up.weight += 1
																								lora_attn_procs[name].to_v_lora.up.weight += 1
																								lora_attn_procs[name].to_out_lora.up.weight += 1

						return lora_attn_procs





@skip_mps
class      snake_case_       (		lowerCAmelCase__     ,	lowerCAmelCase__     ,	unittest.TestCase ):







			'''simple docstring'''

			SCREAMING_SNAKE_CASE					:       Union[str, Any]						      =    UNetaDConditionModel
			SCREAMING_SNAKE_CASE					:       Optional[int]						      =    """sample"""







			@property
			def     snake_case__(					self			:    str							)				->Dict:
									snake_case_       =     4
									snake_case_       =     4
									snake_case_       =     4
									snake_case_       =     (3_2, 3_2)

									snake_case_       =     floats_tensor((batch_size, num_channels, num_frames) + sizes							).to(UpperCamelCase__							)
									snake_case_       =     torch.tensor([1_0]							).to(UpperCamelCase__							)
									snake_case_       =     floats_tensor((batch_size, 4, 3_2)							).to(UpperCamelCase__							)

									return {"sample": noise, "timestep": time_step, "encoder_hidden_states": encoder_hidden_states}







			@property
			def     snake_case__(					self			:    Dict							)				->Optional[int]:
									return (4, 4, 3_2, 3_2)







			@property
			def     snake_case__(					self			:    int							)				->Any:
									return (4, 4, 3_2, 3_2)







			def     snake_case__(					self			:    List[str]							)				->int:
									snake_case_       =     {
									    "block_out_channels": (3_2, 6_4),
									    "down_block_types": (
									        "CrossAttnDownBlock3D",
									        "DownBlock3D",
									    ),
									    "up_block_types": ("UpBlock3D", "CrossAttnUpBlock3D"),
									    "cross_attention_dim": 3_2,
									    "attention_head_dim": 8,
									    "out_channels": 4,
									    "in_channels": 4,
									    "layers_per_block": 1,
									    "sample_size": 3_2,
									}
									snake_case_       =     self.dummy_input
									return init_dict, inputs_dict







			@unittest.skipIf(
			    torch_device != '''cuda''' or not is_xformers_available()							,       reason='''XFormers attention is only available with CUDA and `xformers` installed'''							,       )
			def     snake_case__(					self			:    Optional[int]							)				->Union[str, Any]:
									snake_case_       =     self.prepare_init_args_and_inputs_for_common()
									snake_case_       =     self.model_class(**UpperCamelCase__							)

									model.enable_xformers_memory_efficient_attention()

									assert (
									    model.mid_block.attentions[0].transformer_blocks[0].attna.processor.__class__.__name__
									    == "XFormersAttnProcessor"
									), "xformers is not enabled"







			def     snake_case__(					self			:    Union[str, Any]							)				->Optional[int]:
									snake_case_       =     self.prepare_init_args_and_inputs_for_common()

									snake_case_       =     3_2

									snake_case_       =     self.model_class(**UpperCamelCase__							)
									model.to(UpperCamelCase__							)
									model.eval()

									with torch.no_grad():
															snake_case_       =     model(**UpperCamelCase__							)

															if isinstance(UpperCamelCase__							,       UpperCamelCase__							):
																					snake_case_       =     output.sample

									self.assertIsNotNone(UpperCamelCase__							)
									snake_case_       =     inputs_dict["sample"].shape
									self.assertEqual(output.shape							,       UpperCamelCase__							,       '''Input and output shapes do not match'''							)







			def     snake_case__(					self			:    Dict							)				->List[str]:
									snake_case_       =     self.prepare_init_args_and_inputs_for_common()
									snake_case_       =     self.model_class(**UpperCamelCase__							)
									model.to(UpperCamelCase__							)
									model.eval()

									with torch.no_grad():
															# Warmup pass when using mps (see #372)
															if torch_device == "mps" and isinstance(UpperCamelCase__							,       UpperCamelCase__							):
																					model(**self.dummy_input							)

															snake_case_       =     model(**UpperCamelCase__							)
															if isinstance(UpperCamelCase__							,       UpperCamelCase__							):
																					snake_case_       =     first.sample

															snake_case_       =     model(**UpperCamelCase__							)
															if isinstance(UpperCamelCase__							,       UpperCamelCase__							):
																					snake_case_       =     second.sample

									snake_case_       =     first.cpu().numpy()
									snake_case_       =     second.cpu().numpy()
									snake_case_       =     out_a[~np.isnan(UpperCamelCase__							)]
									snake_case_       =     out_a[~np.isnan(UpperCamelCase__							)]
									snake_case_       =     np.amax(np.abs(out_a - out_a							)							)
									self.assertLessEqual(UpperCamelCase__							,       1e-5							)







			def     snake_case__(					self			:    Tuple							)				->int:
									snake_case_       =     self.prepare_init_args_and_inputs_for_common()

									snake_case_       =     8

									snake_case_       =     self.model_class(**UpperCamelCase__							)
									model.to(UpperCamelCase__							)
									model.eval()

									model.set_attention_slice('''auto'''							)
									with torch.no_grad():
															snake_case_       =     model(**UpperCamelCase__							)
									assert output is not None

									model.set_attention_slice('''max'''							)
									with torch.no_grad():
															snake_case_       =     model(**UpperCamelCase__							)
									assert output is not None

									model.set_attention_slice(2							)
									with torch.no_grad():
															snake_case_       =     model(**UpperCamelCase__							)
									assert output is not None







			def     snake_case__(					self			:    List[str]							)				->str:
									snake_case_       =     self.prepare_init_args_and_inputs_for_common()

									snake_case_       =     8

									snake_case_       =     self.model_class(**UpperCamelCase__							)
									model.to(UpperCamelCase__							)

									with torch.no_grad():
															snake_case_       =     model(**UpperCamelCase__							).sample

									snake_case_       =     create_lora_layers(UpperCamelCase__							)

									# make sure we can set a list of attention processors
									model.set_attn_processor(UpperCamelCase__							)
									model.to(UpperCamelCase__							)

									# test that attn processors can be set to itself
									model.set_attn_processor(model.attn_processors							)

									with torch.no_grad():
															snake_case_       =     model(**UpperCamelCase__							,       cross_attention_kwargs={'''scale''': 0.0}							).sample
															snake_case_       =     model(**UpperCamelCase__							,       cross_attention_kwargs={'''scale''': 0.5}							).sample
															snake_case_       =     model(**UpperCamelCase__							,       cross_attention_kwargs={'''scale''': 0.5}							).sample

									assert (samplea - samplea).abs().max() < 3e-3
									assert (samplea - samplea).abs().max() < 3e-3

									# sample 2 and sample 3 should be different
									assert (samplea - samplea).abs().max() > 3e-3







			def     snake_case__(					self			:    Tuple							)				->Any:
									snake_case_       =     self.prepare_init_args_and_inputs_for_common()

									snake_case_       =     8

									torch.manual_seed(0							)
									snake_case_       =     self.model_class(**UpperCamelCase__							)
									model.to(UpperCamelCase__							)

									with torch.no_grad():
															snake_case_       =     model(**UpperCamelCase__							).sample

									snake_case_       =     create_lora_layers(UpperCamelCase__							)
									model.set_attn_processor(UpperCamelCase__							)

									with torch.no_grad():
															snake_case_       =     model(**UpperCamelCase__							,       cross_attention_kwargs={'''scale''': 0.5}							).sample

									with tempfile.TemporaryDirectory() as tmpdirname:
															model.save_attn_procs(UpperCamelCase__							)
															self.assertTrue(os.path.isfile(os.path.join(UpperCamelCase__							,       '''pytorch_lora_weights.bin'''							)							)							)
															torch.manual_seed(0							)
															snake_case_       =     self.model_class(**UpperCamelCase__							)
															new_model.to(UpperCamelCase__							)
															new_model.load_attn_procs(UpperCamelCase__							)

									with torch.no_grad():
															snake_case_       =     new_model(**UpperCamelCase__							,       cross_attention_kwargs={'''scale''': 0.5}							).sample

									assert (sample - new_sample).abs().max() < 1e-3

									# LoRA and no LoRA should NOT be the same
									assert (sample - old_sample).abs().max() > 1e-4







			def     snake_case__(					self			:    Any							)				->List[str]:
									snake_case_       =     self.prepare_init_args_and_inputs_for_common()

									snake_case_       =     8

									torch.manual_seed(0							)
									snake_case_       =     self.model_class(**UpperCamelCase__							)
									model.to(UpperCamelCase__							)

									with torch.no_grad():
															snake_case_       =     model(**UpperCamelCase__							).sample

									snake_case_       =     create_lora_layers(UpperCamelCase__							)
									model.set_attn_processor(UpperCamelCase__							)

									with torch.no_grad():
															snake_case_       =     model(**UpperCamelCase__							,       cross_attention_kwargs={'''scale''': 0.5}							).sample

									with tempfile.TemporaryDirectory() as tmpdirname:
															model.save_attn_procs(UpperCamelCase__							,       safe_serialization=UpperCamelCase__							)
															self.assertTrue(os.path.isfile(os.path.join(UpperCamelCase__							,       '''pytorch_lora_weights.safetensors'''							)							)							)
															torch.manual_seed(0							)
															snake_case_       =     self.model_class(**UpperCamelCase__							)
															new_model.to(UpperCamelCase__							)
															new_model.load_attn_procs(UpperCamelCase__							)

									with torch.no_grad():
															snake_case_       =     new_model(**UpperCamelCase__							,       cross_attention_kwargs={'''scale''': 0.5}							).sample

									assert (sample - new_sample).abs().max() < 3e-3

									# LoRA and no LoRA should NOT be the same
									assert (sample - old_sample).abs().max() > 1e-4







			def     snake_case__(					self			:    Dict							)				->Any:
									# enable deterministic behavior for gradient checkpointing
									snake_case_       =     self.prepare_init_args_and_inputs_for_common()

									snake_case_       =     8

									torch.manual_seed(0							)
									snake_case_       =     self.model_class(**UpperCamelCase__							)
									model.to(UpperCamelCase__							)

									snake_case_       =     create_lora_layers(UpperCamelCase__							,       mock_weights=UpperCamelCase__							)
									model.set_attn_processor(UpperCamelCase__							)
									# Saving as torch, properly reloads with directly filename
									with tempfile.TemporaryDirectory() as tmpdirname:
															model.save_attn_procs(UpperCamelCase__							)
															self.assertTrue(os.path.isfile(os.path.join(UpperCamelCase__							,       '''pytorch_lora_weights.bin'''							)							)							)
															torch.manual_seed(0							)
															snake_case_       =     self.model_class(**UpperCamelCase__							)
															new_model.to(UpperCamelCase__							)
															new_model.load_attn_procs(UpperCamelCase__							,       weight_name='''pytorch_lora_weights.bin'''							)







			def     snake_case__(					self			:    Optional[int]							)				->int:
									snake_case_       =     self.prepare_init_args_and_inputs_for_common()

									snake_case_       =     8

									torch.manual_seed(0							)
									snake_case_       =     self.model_class(**UpperCamelCase__							)
									model.to(UpperCamelCase__							)

									snake_case_       =     create_lora_layers(UpperCamelCase__							,       mock_weights=UpperCamelCase__							)
									model.set_attn_processor(UpperCamelCase__							)
									# Saving as torch, properly reloads with directly filename
									with tempfile.TemporaryDirectory() as tmpdirname:
															model.save_attn_procs(UpperCamelCase__							)
															self.assertTrue(os.path.isfile(os.path.join(UpperCamelCase__							,       '''pytorch_lora_weights.bin'''							)							)							)
															torch.manual_seed(0							)
															snake_case_       =     self.model_class(**UpperCamelCase__							)
															new_model.to(UpperCamelCase__							)
															with self.assertRaises(UpperCamelCase__							) as e:
																					new_model.load_attn_procs(UpperCamelCase__							,       use_safetensors=UpperCamelCase__							)
															self.assertIn('''Error no file named pytorch_lora_weights.safetensors'''							,       str(e.exception							)							)







			def     snake_case__(					self			:    str							)				->int:
									snake_case_       =     self.prepare_init_args_and_inputs_for_common()

									snake_case_       =     8

									torch.manual_seed(0							)
									snake_case_       =     self.model_class(**UpperCamelCase__							)
									model.to(UpperCamelCase__							)

									with torch.no_grad():
															snake_case_       =     model(**UpperCamelCase__							).sample

									snake_case_       =     create_lora_layers(UpperCamelCase__							)
									model.set_attn_processor(UpperCamelCase__							)

									with torch.no_grad():
															snake_case_       =     model(**UpperCamelCase__							,       cross_attention_kwargs={'''scale''': 0.0}							).sample

									model.set_attn_processor(AttnProcessor()							)

									with torch.no_grad():
															snake_case_       =     model(**UpperCamelCase__							).sample

									assert (sample - new_sample).abs().max() < 1e-4
									assert (sample - old_sample).abs().max() < 3e-3







			@unittest.skipIf(
			    torch_device != '''cuda''' or not is_xformers_available()							,       reason='''XFormers attention is only available with CUDA and `xformers` installed'''							,       )
			def     snake_case__(					self			:    Any							)				->List[Any]:
									# enable deterministic behavior for gradient checkpointing
									snake_case_       =     self.prepare_init_args_and_inputs_for_common()

									snake_case_       =     4

									torch.manual_seed(0							)
									snake_case_       =     self.model_class(**UpperCamelCase__							)
									model.to(UpperCamelCase__							)
									snake_case_       =     create_lora_layers(UpperCamelCase__							)
									model.set_attn_processor(UpperCamelCase__							)

									# default
									with torch.no_grad():
															snake_case_       =     model(**UpperCamelCase__							).sample

															model.enable_xformers_memory_efficient_attention()
															snake_case_       =     model(**UpperCamelCase__							).sample

															model.disable_xformers_memory_efficient_attention()
															snake_case_       =     model(**UpperCamelCase__							).sample

									assert (sample - on_sample).abs().max() < 1e-4
									assert (sample - off_sample).abs().max() < 1e-4







			def     snake_case__(					self			:    Optional[int]							)				->Union[str, Any]:
									snake_case_       =     self.prepare_init_args_and_inputs_for_common()
									snake_case_       =     3_2

									snake_case_       =     self.model_class(**UpperCamelCase__							)
									model.to(UpperCamelCase__							)
									model.eval()

									with torch.no_grad():
															snake_case_       =     model(**UpperCamelCase__							)[0]

									model.enable_forward_chunking()
									with torch.no_grad():
															snake_case_       =     model(**UpperCamelCase__							)[0]

									self.assertEqual(output.shape							,       output_a.shape							,       '''Shape doesn\'t match'''							)
									assert np.abs(output.cpu() - output_a.cpu()							).max() < 1e-2


# (todo: sayakpaul) implement SLOW tests.


import json
from typing import List, Optional, Tuple

from tokenizers import normalizers

from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import logging
from .tokenization_splinter import SplinterTokenizer


lowerCAmelCase_	     =     logging.get_logger(__name__)

lowerCAmelCase_	     =     {'''vocab_file''': '''vocab.txt'''}

lowerCAmelCase_	     =     {
    '''vocab_file''': {
        '''tau/splinter-base''': '''https://huggingface.co/tau/splinter-base/resolve/main/vocab.txt''',
        '''tau/splinter-base-qass''': '''https://huggingface.co/tau/splinter-base-qass/resolve/main/vocab.txt''',
        '''tau/splinter-large''': '''https://huggingface.co/tau/splinter-large/resolve/main/vocab.txt''',
        '''tau/splinter-large-qass''': '''https://huggingface.co/tau/splinter-large-qass/resolve/main/vocab.txt''',
    }
}

lowerCAmelCase_	     =     {
    '''tau/splinter-base''': 5_12,
    '''tau/splinter-base-qass''': 5_12,
    '''tau/splinter-large''': 5_12,
    '''tau/splinter-large-qass''': 5_12,
}

lowerCAmelCase_	     =     {
    '''tau/splinter-base''': {'''do_lower_case''': False},
    '''tau/splinter-base-qass''': {'''do_lower_case''': False},
    '''tau/splinter-large''': {'''do_lower_case''': False},
    '''tau/splinter-large-qass''': {'''do_lower_case''': False},
}





class      snake_case_       (		__A ):







			'''simple docstring'''

			SCREAMING_SNAKE_CASE					:       List[str]						      =    VOCAB_FILES_NAMES
			SCREAMING_SNAKE_CASE					:       List[str]						      =    PRETRAINED_VOCAB_FILES_MAP
			SCREAMING_SNAKE_CASE					:       Optional[int]						      =    PRETRAINED_INIT_CONFIGURATION
			SCREAMING_SNAKE_CASE					:       List[str]						      =    PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
			SCREAMING_SNAKE_CASE					:       Any						      =    SplinterTokenizer







			def __init__(					self			:    List[str]							,       _UpperCamelCase			:    Dict=None							,       _UpperCamelCase			:    Union[str, Any]=None							,       _UpperCamelCase			:    Union[str, Any]=True							,       _UpperCamelCase			:    List[Any]="[UNK]"							,       _UpperCamelCase			:    List[str]="[SEP]"							,       _UpperCamelCase			:    Optional[int]="[PAD]"							,       _UpperCamelCase			:    Tuple="[CLS]"							,       _UpperCamelCase			:    Optional[int]="[MASK]"							,       _UpperCamelCase			:    str="[QUESTION]"							,       _UpperCamelCase			:    Optional[Any]=True							,       _UpperCamelCase			:    Dict=None							,       **_UpperCamelCase			:    Optional[int]							,       )				->Union[str, Any]:
									super().__init__(
									    _UpperCamelCase							,       tokenizer_file=_UpperCamelCase							,       do_lower_case=_UpperCamelCase							,       unk_token=_UpperCamelCase							,       sep_token=_UpperCamelCase							,       pad_token=_UpperCamelCase							,       cls_token=_UpperCamelCase							,       mask_token=_UpperCamelCase							,       tokenize_chinese_chars=_UpperCamelCase							,       strip_accents=_UpperCamelCase							,       additional_special_tokens=(question_token,)							,       **_UpperCamelCase							,       )

									snake_case_       =     json.loads(self.backend_tokenizer.normalizer.__getstate__()							)
									if (
									    pre_tok_state.get('''lowercase'''							,       _UpperCamelCase							) != do_lower_case
									    or pre_tok_state.get('''strip_accents'''							,       _UpperCamelCase							) != strip_accents
									):
															snake_case_       =     getattr(_UpperCamelCase							,       pre_tok_state.pop('''type'''							)							)
															snake_case_       =     do_lower_case
															snake_case_       =     strip_accents
															snake_case_       =     pre_tok_class(**_UpperCamelCase							)

									snake_case_       =     do_lower_case







			@property
			def     snake_case__(					self			:    Optional[int]							)				->int:
									return self.convert_tokens_to_ids(self.question_token							)







			def     snake_case__(					self			:    Dict							,       _UpperCamelCase			:    List[int]							,       _UpperCamelCase			:    Optional[List[int]] = None							)				->List[int]:
									if token_ids_a is None:
															return [self.cls_token_id] + token_ids_a + [self.sep_token_id]

									snake_case_       =     [self.cls_token_id]
									snake_case_       =     [self.sep_token_id]
									snake_case_       =     [self.question_token_id] + [self.convert_tokens_to_ids('''.'''							)]
									if self.padding_side == "right":
															# Input is question-then-context
															return cls + token_ids_a + question_suffix + sep + token_ids_a + sep
									else:
															# Input is context-then-question
															return cls + token_ids_a + sep + token_ids_a + question_suffix + sep







			def     snake_case__(					self			:    str							,       _UpperCamelCase			:    List[int]							,       _UpperCamelCase			:    Optional[List[int]] = None							)				->List[int]:
									snake_case_       =     [self.sep_token_id]
									snake_case_       =     [self.cls_token_id]
									snake_case_       =     [self.question_token_id] + [self.convert_tokens_to_ids('''.'''							)]
									if token_ids_a is None:
															return len(cls + token_ids_a + sep							) * [0]

									if self.padding_side == "right":
															# Input is question-then-context
															return len(cls + token_ids_a + question_suffix + sep							) * [0] + len(token_ids_a + sep							) * [1]
									else:
															# Input is context-then-question
															return len(cls + token_ids_a + sep							) * [0] + len(token_ids_a + question_suffix + sep							) * [1]







			def     snake_case__(					self			:    Union[str, Any]							,       _UpperCamelCase			:    str							,       _UpperCamelCase			:    Optional[str] = None							)				->Tuple[str]:
									snake_case_       =     self._tokenizer.model.save(_UpperCamelCase							,       name=_UpperCamelCase							)
									return tuple(_UpperCamelCase							)