




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       (    UpperCamelCase__	,				UpperCamelCase__	,				UpperCamelCase__	,				unittest.TestCase    ):




    """simple docstring"""

    snake_case__					=      StableDiffusionModelEditingPipeline
    snake_case__					=      TEXT_TO_IMAGE_PARAMS
    snake_case__					=      TEXT_TO_IMAGE_BATCH_PARAMS
    snake_case__					=      TEXT_TO_IMAGE_IMAGE_PARAMS
    snake_case__					=      TEXT_TO_IMAGE_IMAGE_PARAMS


    def    a      (       self     :     List[Any]   )      ->       Dict:
         torch.manual_seed(0   )
         lowerCAmelCase__		= UNetaDConditionModel(
             block_out_channels=(32, 64)   ,	layers_per_block=2   ,	sample_size=32   ,	in_channels=4   ,	out_channels=4   ,	down_block_types=("DownBlock2D", "CrossAttnDownBlock2D")   ,	up_block_types=("CrossAttnUpBlock2D", "UpBlock2D")   ,	cross_attention_dim=32   ,	)
         lowerCAmelCase__		= DDIMScheduler()
         torch.manual_seed(0   )
         lowerCAmelCase__		= AutoencoderKL(
             block_out_channels=[32, 64]   ,	in_channels=3   ,	out_channels=3   ,	down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"]   ,	up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"]   ,	latent_channels=4   ,	)
         torch.manual_seed(0   )
         lowerCAmelCase__		= CLIPTextConfig(
             bos_token_id=0   ,	eos_token_id=2   ,	hidden_size=32   ,	intermediate_size=37   ,	layer_norm_eps=1e-0_5   ,	num_attention_heads=4   ,	num_hidden_layers=5   ,	pad_token_id=1   ,	vocab_size=1_000   ,	)
         lowerCAmelCase__		= CLIPTextModel(SCREAMING_SNAKE_CASE__   )
         lowerCAmelCase__		= CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"   )

         lowerCAmelCase__		= {
             "unet": unet,
             "scheduler": scheduler,
             "vae": vae,
             "text_encoder": text_encoder,
             "tokenizer": tokenizer,
             "safety_checker": None,
             "feature_extractor": None,
         }
         return components


    def    a      (       self     :     Any   ,	SCREAMING_SNAKE_CASE__     :     Any   ,	SCREAMING_SNAKE_CASE__     :     Tuple=0   )      ->       Optional[Any]:
         lowerCAmelCase__		= torch.manual_seed(SCREAMING_SNAKE_CASE__   )
         lowerCAmelCase__		= {
             "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    a      (       self     :     Optional[int]   )      ->       Optional[int]:
         lowerCAmelCase__		= "cpu"  # ensure determinism for the device-dependent torch.Generator
         lowerCAmelCase__		= self.get_dummy_components()
         lowerCAmelCase__		= StableDiffusionModelEditingPipeline(**SCREAMING_SNAKE_CASE__   )
         lowerCAmelCase__		= sd_pipe.to(SCREAMING_SNAKE_CASE__   )
         sd_pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__   )

         lowerCAmelCase__		= self.get_dummy_inputs(SCREAMING_SNAKE_CASE__   )
         lowerCAmelCase__		= sd_pipe(**SCREAMING_SNAKE_CASE__   ).images
         lowerCAmelCase__		= image[0, -3:, -3:, -1]
         assert image.shape == (1, 64, 64, 3)

         lowerCAmelCase__		= np.array([0.4_755, 0.5_132, 0.4_976, 0.3_904, 0.3_554, 0.4_765, 0.5_139, 0.5_158, 0.4_889]   )

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


    def    a      (       self     :     Union[str, Any]   )      ->       Optional[int]:
         lowerCAmelCase__		= "cpu"  # ensure determinism for the device-dependent torch.Generator
         lowerCAmelCase__		= self.get_dummy_components()
         lowerCAmelCase__		= StableDiffusionModelEditingPipeline(**SCREAMING_SNAKE_CASE__   )
         lowerCAmelCase__		= sd_pipe.to(SCREAMING_SNAKE_CASE__   )
         sd_pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__   )

         lowerCAmelCase__		= self.get_dummy_inputs(SCREAMING_SNAKE_CASE__   )
         lowerCAmelCase__		= "french fries"
         lowerCAmelCase__		= sd_pipe(**SCREAMING_SNAKE_CASE__   ,	negative_prompt=SCREAMING_SNAKE_CASE__   )
         lowerCAmelCase__		= output.images
         lowerCAmelCase__		= image[0, -3:, -3:, -1]

         assert image.shape == (1, 64, 64, 3)

         lowerCAmelCase__		= np.array([0.4_992, 0.5_101, 0.5_004, 0.3_949, 0.3_604, 0.4_735, 0.5_216, 0.5_204, 0.4_913]   )

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


    def    a      (       self     :     str   )      ->       List[Any]:
         lowerCAmelCase__		= "cpu"  # ensure determinism for the device-dependent torch.Generator
         lowerCAmelCase__		= self.get_dummy_components()
         lowerCAmelCase__		= EulerAncestralDiscreteScheduler(
             beta_start=0.00_085   ,	beta_end=0.012   ,	beta_schedule="scaled_linear"   )
         lowerCAmelCase__		= StableDiffusionModelEditingPipeline(**SCREAMING_SNAKE_CASE__   )
         lowerCAmelCase__		= sd_pipe.to(SCREAMING_SNAKE_CASE__   )
         sd_pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__   )

         lowerCAmelCase__		= self.get_dummy_inputs(SCREAMING_SNAKE_CASE__   )
         lowerCAmelCase__		= sd_pipe(**SCREAMING_SNAKE_CASE__   ).images
         lowerCAmelCase__		= image[0, -3:, -3:, -1]

         assert image.shape == (1, 64, 64, 3)

         lowerCAmelCase__		= np.array([0.4_747, 0.5_372, 0.4_779, 0.4_982, 0.5_543, 0.4_816, 0.5_238, 0.4_904, 0.5_027]   )

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


    def    a      (       self     :     Union[str, Any]   )      ->       Optional[Any]:
         lowerCAmelCase__		= "cpu"  # ensure determinism for the device-dependent torch.Generator
         lowerCAmelCase__		= self.get_dummy_components()
         lowerCAmelCase__		= PNDMScheduler()
         lowerCAmelCase__		= StableDiffusionModelEditingPipeline(**SCREAMING_SNAKE_CASE__   )
         lowerCAmelCase__		= sd_pipe.to(SCREAMING_SNAKE_CASE__   )
         sd_pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__   )

         lowerCAmelCase__		= self.get_dummy_inputs(SCREAMING_SNAKE_CASE__   )
         # the pipeline does not expect pndm so test if it raises error.
         with self.assertRaises(SCREAMING_SNAKE_CASE__   ):
              lowerCAmelCase__		= sd_pipe(**SCREAMING_SNAKE_CASE__   ).images


    def    a      (       self     :     Any   )      ->       Tuple:
         super().test_inference_batch_single_identical(expected_max_diff=5e-3   )


    def    a      (       self     :     Optional[Any]   )      ->       Dict:
         super().test_attention_slicing_forward_pass(expected_max_diff=5e-3   )






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




    """simple docstring"""

    def    a      (       self     :     str   )      ->       Any:
         super().tearDown()
         gc.collect()
         torch.cuda.empty_cache()


    def    a      (       self     :     Tuple   ,	SCREAMING_SNAKE_CASE__     :     Any=0   )      ->       Dict:
         lowerCAmelCase__		= torch.manual_seed(SCREAMING_SNAKE_CASE__   )
         lowerCAmelCase__		= {
             "prompt": "A field of roses",
             "generator": generator,
             "num_inference_steps": 3,
             "guidance_scale": 7.5,
             "output_type": "numpy",
         }
         return inputs


    def    a      (       self     :     str   )      ->       List[Any]:
         lowerCAmelCase__		= "CompVis/stable-diffusion-v1-4"
         lowerCAmelCase__		= StableDiffusionModelEditingPipeline.from_pretrained(SCREAMING_SNAKE_CASE__   ,	safety_checker=SCREAMING_SNAKE_CASE__   )
         pipe.to(SCREAMING_SNAKE_CASE__   )
         pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__   )
         pipe.enable_attention_slicing()

         lowerCAmelCase__		= self.get_inputs()
         lowerCAmelCase__		= pipe(**SCREAMING_SNAKE_CASE__   ).images
         lowerCAmelCase__		= image[0, -3:, -3:, -1].flatten()

         assert image.shape == (1, 512, 512, 3)

         lowerCAmelCase__		= np.array(
             [0.6_749_496, 0.6_386_453, 0.51_443_267, 0.66_094_905, 0.61_921_215, 0.5_491_332, 0.5_744_417, 0.58_075_106, 0.5_174_658]   )

         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"   )

         lowerCAmelCase__		= pipe(**SCREAMING_SNAKE_CASE__   ).images
         lowerCAmelCase__		= image[0, -3:, -3:, -1].flatten()

         assert image.shape == (1, 512, 512, 3)

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


    def    a      (       self     :     int   )      ->       Tuple:
         torch.cuda.empty_cache()
         torch.cuda.reset_max_memory_allocated()
         torch.cuda.reset_peak_memory_stats()

         lowerCAmelCase__		= "CompVis/stable-diffusion-v1-4"
         lowerCAmelCase__		= DDIMScheduler.from_pretrained(SCREAMING_SNAKE_CASE__   ,	subfolder="scheduler"   )
         lowerCAmelCase__		= StableDiffusionModelEditingPipeline.from_pretrained(
             SCREAMING_SNAKE_CASE__   ,	scheduler=SCREAMING_SNAKE_CASE__   ,	safety_checker=SCREAMING_SNAKE_CASE__   )
         lowerCAmelCase__		= pipe.to(SCREAMING_SNAKE_CASE__   )
         pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__   )
         pipe.enable_attention_slicing(1   )
         pipe.enable_sequential_cpu_offload()

         lowerCAmelCase__		= self.get_inputs()
         lowerCAmelCase__		= pipe(**SCREAMING_SNAKE_CASE__   )

         lowerCAmelCase__		= torch.cuda.max_memory_allocated()
         # make sure that less than 4.4 GB is allocated
         assert mem_bytes < 4.4 * 10**9






import os
import tempfile
import unittest
from typing import List

from transformers.models.esm.tokenization_esm import VOCAB_FILES_NAMES, EsmTokenizer
from transformers.testing_utils import require_tokenizers
from transformers.tokenization_utils import PreTrainedTokenizer
from transformers.tokenization_utils_base import PreTrainedTokenizerBase





@require_tokenizers
class  __lowerCamelCase       (    unittest.TestCase    ):




    """simple docstring"""

    snake_case__					=      EsmTokenizer


    def    a      (       self     :     str   )      ->       Tuple:
         super().setUp()
         lowerCAmelCase__		= tempfile.mkdtemp()
         # fmt: off
         lowerCAmelCase__		= ["<cls>", "<pad>", "<eos>", "<unk>", "L", "A", "G", "V", "S", "E", "R", "T", "I", "D", "P", "K", "Q", "N", "F", "Y", "M", "H", "W", "C", "X", "B", "U", "Z", "O", ".", "-", "<null_1>", "<mask>"]  # noqa: E501
         # fmt: on
         lowerCAmelCase__		= 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    a      (       self     :     Dict   ,	**SCREAMING_SNAKE_CASE__     :     Optional[Any]   )      ->       List[PreTrainedTokenizerBase]:
         return [self.get_tokenizer(**SCREAMING_SNAKE_CASE__   )]


    def    a      (       self     :     List[str]   ,	**SCREAMING_SNAKE_CASE__     :     str   )      ->       PreTrainedTokenizer:
         return self.tokenizer_class.from_pretrained(self.tmpdirname   ,	**SCREAMING_SNAKE_CASE__   )


    def    a      (       self     :     str   )      ->       Tuple:
         lowerCAmelCase__		= self.tokenizer_class(self.vocab_file   )

         lowerCAmelCase__		= tokenizer.tokenize("LAGVS"   )
         self.assertListEqual(SCREAMING_SNAKE_CASE__   ,	["L", "A", "G", "V", "S"]   )
         self.assertListEqual(tokenizer.convert_tokens_to_ids(SCREAMING_SNAKE_CASE__   )   ,	[4, 5, 6, 7, 8]   )


    def    a      (       self     :     int   )      ->       Union[str, Any]:
         lowerCAmelCase__		= self.tokenizer_class(self.vocab_file   )

         lowerCAmelCase__		= "LAGVS"
         self.assertListEqual(tokenizer.encode(SCREAMING_SNAKE_CASE__   )   ,	[0, 4, 5, 6, 7, 8, 2]   )


    def    a      (       self     :     Union[str, Any]   )      ->       str:
         lowerCAmelCase__		= self.tokenizer_class(self.vocab_file   )

         lowerCAmelCase__		= ["LAGVS", "WCB"]
         lowerCAmelCase__		= tokenizer(SCREAMING_SNAKE_CASE__   ,	padding=SCREAMING_SNAKE_CASE__   )["input_ids"]

         self.assertListEqual(SCREAMING_SNAKE_CASE__   ,	[[0, 4, 5, 6, 7, 8, 2], [0, 22, 23, 25, 2]]   )


    def    a      (       self     :     Any   )      ->       Union[str, Any]:
         lowerCAmelCase__		= self.tokenizer_class(self.vocab_file   )

         lowerCAmelCase__		= ["LAGVS", "WCB"]
         lowerCAmelCase__		= tokenizer(SCREAMING_SNAKE_CASE__   ,	padding=SCREAMING_SNAKE_CASE__   )["input_ids"]

         self.assertListEqual(SCREAMING_SNAKE_CASE__   ,	[[0, 4, 5, 6, 7, 8, 2], [0, 22, 23, 25, 2, 1, 1]]   )


    def    a      (       self     :     str   )      ->       Optional[Any]:
         lowerCAmelCase__		= self.get_tokenizers(fast=SCREAMING_SNAKE_CASE__   )
         for tokenizer in tokenizers:
              with self.subTest(f'{tokenizer.__class__.__name__}'   ):
                   lowerCAmelCase__		= "<unk>"
                   lowerCAmelCase__		= "<mask>"

                   lowerCAmelCase__		= tokenizer.tokenize(SCREAMING_SNAKE_CASE__   )
                   lowerCAmelCase__		= tokenizer.tokenize(SCREAMING_SNAKE_CASE__   )

                   self.assertEqual(len(SCREAMING_SNAKE_CASE__   )   ,	1   )
                   self.assertEqual(len(SCREAMING_SNAKE_CASE__   )   ,	1   )
                   self.assertEqual(token_a[0]   ,	SCREAMING_SNAKE_CASE__   )
                   self.assertEqual(token_a[0]   ,	SCREAMING_SNAKE_CASE__   )
