


import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    EulerAncestralDiscreteScheduler,
    PNDMScheduler,
    StableDiffusionModelEditingPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, skip_mps

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()







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




     '''simple docstring'''
     __SCREAMING_SNAKE_CASE        =			StableDiffusionModelEditingPipeline
     __SCREAMING_SNAKE_CASE        =			TEXT_TO_IMAGE_PARAMS
     __SCREAMING_SNAKE_CASE        =			TEXT_TO_IMAGE_BATCH_PARAMS
     __SCREAMING_SNAKE_CASE        =			TEXT_TO_IMAGE_IMAGE_PARAMS
     __SCREAMING_SNAKE_CASE        =			TEXT_TO_IMAGE_IMAGE_PARAMS


     def 			_lowerCamelCase			(						self)      ->    Optional[int]:
            torch.manual_seed(0)
            _A				:	str       =			UNetaDConditionModel(
                block_out_channels=(3_2, 6_4)			,						layers_per_block=2			,						sample_size=3_2			,						in_channels=4			,						out_channels=4			,						down_block_types=("DownBlock2D", "CrossAttnDownBlock2D")			,						up_block_types=("CrossAttnUpBlock2D", "UpBlock2D")			,						cross_attention_dim=3_2			,						)
            _A				:	Union[str, Any]       =			DDIMScheduler()
            torch.manual_seed(0)
            _A				:	Any       =			AutoencoderKL(
                block_out_channels=[3_2, 6_4]			,						in_channels=3			,						out_channels=3			,						down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"]			,						up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"]			,						latent_channels=4			,						)
            torch.manual_seed(0)
            _A				:	int       =			CLIPTextConfig(
                bos_token_id=0			,						eos_token_id=2			,						hidden_size=3_2			,						intermediate_size=3_7			,						layer_norm_eps=1e-05			,						num_attention_heads=4			,						num_hidden_layers=5			,						pad_token_id=1			,						vocab_size=1_0_0_0			,						)
            _A				:	Union[str, Any]       =			CLIPTextModel(__lowerCamelCase)
            _A				:	Tuple       =			CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")

            _A				:	Optional[Any]       =			{
                "unet": unet,
                "scheduler": scheduler,
                "vae": vae,
                "text_encoder": text_encoder,
                "tokenizer": tokenizer,
                "safety_checker": None,
                "feature_extractor": None,
            }
            return components


     def 			_lowerCamelCase			(						self			,						__lowerCamelCase			,						__lowerCamelCase=0)      ->    int:
            _A				:	Any       =			torch.manual_seed(__lowerCamelCase)
            _A				:	List[str]       =			{
                "prompt": "A field of roses",
                "generator": generator,
                # Setting height and width to None to prevent OOMs on CPU.
                "height": None,
                "width": None,
                "num_inference_steps": 2,
                "guidance_scale": 6.0,
                "output_type": "numpy",
            }
            return inputs


     def 			_lowerCamelCase			(						self)      ->    Optional[Any]:
            _A				:	Optional[Any]       =			"cpu"  # ensure determinism for the device-dependent torch.Generator
            _A				:	List[str]       =			self.get_dummy_components()
            _A				:	int       =			StableDiffusionModelEditingPipeline(**__lowerCamelCase)
            _A				:	Tuple       =			sd_pipe.to(__lowerCamelCase)
            sd_pipe.set_progress_bar_config(disable=__lowerCamelCase)

            _A				:	int       =			self.get_dummy_inputs(__lowerCamelCase)
            _A				:	Tuple       =			sd_pipe(**__lowerCamelCase).images
            _A				:	List[Any]       =			image[0, -3:, -3:, -1]
            assert image.shape == (1, 6_4, 6_4, 3)

            _A				:	Union[str, Any]       =			np.array([0.4_7_5_5, 0.5_1_3_2, 0.4_9_7_6, 0.3_9_0_4, 0.3_5_5_4, 0.4_7_6_5, 0.5_1_3_9, 0.5_1_5_8, 0.4_8_8_9])

            assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2


     def 			_lowerCamelCase			(						self)      ->    List[Any]:
            _A				:	Optional[int]       =			"cpu"  # ensure determinism for the device-dependent torch.Generator
            _A				:	Optional[Any]       =			self.get_dummy_components()
            _A				:	str       =			StableDiffusionModelEditingPipeline(**__lowerCamelCase)
            _A				:	str       =			sd_pipe.to(__lowerCamelCase)
            sd_pipe.set_progress_bar_config(disable=__lowerCamelCase)

            _A				:	List[Any]       =			self.get_dummy_inputs(__lowerCamelCase)
            _A				:	Dict       =			"french fries"
            _A				:	int       =			sd_pipe(**__lowerCamelCase			,						negative_prompt=__lowerCamelCase)
            _A				:	List[str]       =			output.images
            _A				:	List[Any]       =			image[0, -3:, -3:, -1]

            assert image.shape == (1, 6_4, 6_4, 3)

            _A				:	Any       =			np.array([0.4_9_9_2, 0.5_1_0_1, 0.5_0_0_4, 0.3_9_4_9, 0.3_6_0_4, 0.4_7_3_5, 0.5_2_1_6, 0.5_2_0_4, 0.4_9_1_3])

            assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2


     def 			_lowerCamelCase			(						self)      ->    str:
            _A				:	Optional[int]       =			"cpu"  # ensure determinism for the device-dependent torch.Generator
            _A				:	Optional[int]       =			self.get_dummy_components()
            _A				:	Union[str, Any]       =			EulerAncestralDiscreteScheduler(
                beta_start=0.0_0_0_8_5			,						beta_end=0.0_1_2			,						beta_schedule="scaled_linear")
            _A				:	Optional[Any]       =			StableDiffusionModelEditingPipeline(**__lowerCamelCase)
            _A				:	str       =			sd_pipe.to(__lowerCamelCase)
            sd_pipe.set_progress_bar_config(disable=__lowerCamelCase)

            _A				:	Optional[Any]       =			self.get_dummy_inputs(__lowerCamelCase)
            _A				:	Tuple       =			sd_pipe(**__lowerCamelCase).images
            _A				:	List[str]       =			image[0, -3:, -3:, -1]

            assert image.shape == (1, 6_4, 6_4, 3)

            _A				:	List[Any]       =			np.array([0.4_7_4_7, 0.5_3_7_2, 0.4_7_7_9, 0.4_9_8_2, 0.5_5_4_3, 0.4_8_1_6, 0.5_2_3_8, 0.4_9_0_4, 0.5_0_2_7])

            assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2


     def 			_lowerCamelCase			(						self)      ->    Any:
            _A				:	Optional[Any]       =			"cpu"  # ensure determinism for the device-dependent torch.Generator
            _A				:	List[Any]       =			self.get_dummy_components()
            _A				:	Optional[int]       =			PNDMScheduler()
            _A				:	Optional[int]       =			StableDiffusionModelEditingPipeline(**__lowerCamelCase)
            _A				:	Tuple       =			sd_pipe.to(__lowerCamelCase)
            sd_pipe.set_progress_bar_config(disable=__lowerCamelCase)

            _A				:	Any       =			self.get_dummy_inputs(__lowerCamelCase)
            # the pipeline does not expect pndm so test if it raises error.
            with self.assertRaises(__lowerCamelCase):
                   _A				:	Union[str, Any]       =			sd_pipe(**__lowerCamelCase).images


     def 			_lowerCamelCase			(						self)      ->    Optional[Any]:
            super().test_inference_batch_single_identical(expected_max_diff=5e-3)


     def 			_lowerCamelCase			(						self)      ->    Any:
            super().test_attention_slicing_forward_pass(expected_max_diff=5e-3)







@slow
@require_torch_gpu
class 					lowerCAmelCase__	(    unittest.TestCase):




     '''simple docstring'''


     def 			_lowerCamelCase			(						self)      ->    List[str]:
            super().tearDown()
            gc.collect()
            torch.cuda.empty_cache()


     def 			_lowerCamelCase			(						self			,						__lowerCamelCase=0)      ->    Optional[int]:
            _A				:	Any       =			torch.manual_seed(__lowerCamelCase)
            _A				:	Dict       =			{
                "prompt": "A field of roses",
                "generator": generator,
                "num_inference_steps": 3,
                "guidance_scale": 7.5,
                "output_type": "numpy",
            }
            return inputs


     def 			_lowerCamelCase			(						self)      ->    str:
            _A				:	int       =			"CompVis/stable-diffusion-v1-4"
            _A				:	Dict       =			StableDiffusionModelEditingPipeline.from_pretrained(__lowerCamelCase			,						safety_checker=__lowerCamelCase)
            pipe.to(__lowerCamelCase)
            pipe.set_progress_bar_config(disable=__lowerCamelCase)
            pipe.enable_attention_slicing()

            _A				:	List[Any]       =			self.get_inputs()
            _A				:	Any       =			pipe(**__lowerCamelCase).images
            _A				:	List[Any]       =			image[0, -3:, -3:, -1].flatten()

            assert image.shape == (1, 5_1_2, 5_1_2, 3)

            _A				:	List[str]       =			np.array(
                [0.6_7_4_9_4_9_6, 0.6_3_8_6_4_5_3, 0.5_1_4_4_3_2_6_7, 0.6_6_0_9_4_9_0_5, 0.6_1_9_2_1_2_1_5, 0.5_4_9_1_3_3_2, 0.5_7_4_4_4_1_7, 0.5_8_0_7_5_1_0_6, 0.5_1_7_4_6_5_8])

            assert np.abs(expected_slice - image_slice).max() < 1e-2

            # make sure image changes after editing
            pipe.edit_model("A pack of roses"			,						"A pack of blue roses")

            _A				:	List[str]       =			pipe(**__lowerCamelCase).images
            _A				:	str       =			image[0, -3:, -3:, -1].flatten()

            assert image.shape == (1, 5_1_2, 5_1_2, 3)

            assert np.abs(expected_slice - image_slice).max() > 1e-1


     def 			_lowerCamelCase			(						self)      ->    Union[str, Any]:
            torch.cuda.empty_cache()
            torch.cuda.reset_max_memory_allocated()
            torch.cuda.reset_peak_memory_stats()

            _A				:	List[Any]       =			"CompVis/stable-diffusion-v1-4"
            _A				:	Union[str, Any]       =			DDIMScheduler.from_pretrained(__lowerCamelCase			,						subfolder="scheduler")
            _A				:	Any       =			StableDiffusionModelEditingPipeline.from_pretrained(
                __lowerCamelCase			,						scheduler=__lowerCamelCase			,						safety_checker=__lowerCamelCase)
            _A				:	List[Any]       =			pipe.to(__lowerCamelCase)
            pipe.set_progress_bar_config(disable=__lowerCamelCase)
            pipe.enable_attention_slicing(1)
            pipe.enable_sequential_cpu_offload()

            _A				:	str       =			self.get_inputs()
            _A				:	str       =			pipe(**__lowerCamelCase)

            _A				:	int       =			torch.cuda.max_memory_allocated()
            # make sure that less than 4.4 GB is allocated
            assert mem_bytes < 4.4 * 1_0**9




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 					lowerCAmelCase__	(    a):




     '''simple docstring'''
     __SCREAMING_SNAKE_CASE        =			VOCAB_FILES_NAMES
     __SCREAMING_SNAKE_CASE        =			PRETRAINED_VOCAB_FILES_MAP
     __SCREAMING_SNAKE_CASE        =			PRETRAINED_INIT_CONFIGURATION
     __SCREAMING_SNAKE_CASE        =			PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
     __SCREAMING_SNAKE_CASE        =			SplinterTokenizer


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

            _A				:	Tuple       =			json.loads(self.backend_tokenizer.normalizer.__getstate__())
            if (
                pre_tok_state.get("lowercase"			,						__lowerCamelCase) != do_lower_case
                or pre_tok_state.get("strip_accents"			,						__lowerCamelCase) != strip_accents
            ):
                   _A				:	Dict       =			getattr(__lowerCamelCase			,						pre_tok_state.pop("type"))
                   _A				:	List[str]       =			do_lower_case
                   _A				:	Any       =			strip_accents
                   _A				:	int       =			pre_tok_class(**__lowerCamelCase)

            _A				:	int       =			do_lower_case


     @property
     def 			_lowerCamelCase			(						self)      ->    str:
            return self.convert_tokens_to_ids(self.question_token)


     def 			_lowerCamelCase			(						self			,						__lowerCamelCase			,						__lowerCamelCase = None)      ->    List[int]:
            if token_ids_a is None:
                   return [self.cls_token_id] + token_ids_a + [self.sep_token_id]

            _A				:	Tuple       =			[self.cls_token_id]
            _A				:	str       =			[self.sep_token_id]
            _A				:	int       =			[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 			_lowerCamelCase			(						self			,						__lowerCamelCase			,						__lowerCamelCase = None)      ->    List[int]:
            _A				:	str       =			[self.sep_token_id]
            _A				:	int       =			[self.cls_token_id]
            _A				:	Union[str, Any]       =			[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 			_lowerCamelCase			(						self			,						__lowerCamelCase			,						__lowerCamelCase = None)      ->    Tuple[str]:
            _A				:	Tuple       =			self._tokenizer.model.save(__lowerCamelCase			,						name=__lowerCamelCase)
            return tuple(__lowerCamelCase)
