





"""simple docstring"""


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      lowercase_				(     __lowerCAmelCase				,__lowerCAmelCase				,__lowerCAmelCase				,unittest.TestCase							):
       snake_case__				:							Any											=     StableDiffusionModelEditingPipeline
       snake_case__				:							Tuple											=     TEXT_TO_IMAGE_PARAMS
       snake_case__				:							List[Any]											=     TEXT_TO_IMAGE_BATCH_PARAMS
       snake_case__				:							List[Any]											=     TEXT_TO_IMAGE_IMAGE_PARAMS
       snake_case__				:							Union[str, Any]											=     TEXT_TO_IMAGE_IMAGE_PARAMS

       def 		_SCREAMING_SNAKE_CASE(			self     ):






        '''simple docstring'''

        torch.manual_seed(0     )
        snake_case__    							=				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   ,					)
        snake_case__    							=				DDIMScheduler()
        torch.manual_seed(0     )
        snake_case__    							=				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     )
        snake_case__    							=				CLIPTextConfig(
            bos_token_id=0   ,					eos_token_id=2   ,					hidden_size=32   ,					intermediate_size=37   ,					layer_norm_eps=1e-05   ,					num_attention_heads=4   ,					num_hidden_layers=5   ,					pad_token_id=1   ,					vocab_size=1000   ,					)
        snake_case__    							=				CLIPTextModel(__SCREAMING_SNAKE_CASE     )
        snake_case__    							=				CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"     )

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

       def 		_SCREAMING_SNAKE_CASE(			self   ,					__SCREAMING_SNAKE_CASE   ,					__SCREAMING_SNAKE_CASE=0     ):






        '''simple docstring'''

        snake_case__    							=				torch.manual_seed(__SCREAMING_SNAKE_CASE     )
        snake_case__    							=				{
            "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 		_SCREAMING_SNAKE_CASE(			self     ):






        '''simple docstring'''

        snake_case__    							=				"cpu"  # ensure determinism for the device-dependent torch.Generator
        snake_case__    							=				self.get_dummy_components()
        snake_case__    							=				StableDiffusionModelEditingPipeline(**__SCREAMING_SNAKE_CASE     )
        snake_case__    							=				sd_pipe.to(__SCREAMING_SNAKE_CASE     )
        sd_pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE     )

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

        snake_case__    							=				np.array([0.4755, 0.5132, 0.4976, 0.3904, 0.3554, 0.4765, 0.5139, 0.5158, 0.4889]     )

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

       def 		_SCREAMING_SNAKE_CASE(			self     ):






        '''simple docstring'''

        snake_case__    							=				"cpu"  # ensure determinism for the device-dependent torch.Generator
        snake_case__    							=				self.get_dummy_components()
        snake_case__    							=				StableDiffusionModelEditingPipeline(**__SCREAMING_SNAKE_CASE     )
        snake_case__    							=				sd_pipe.to(__SCREAMING_SNAKE_CASE     )
        sd_pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE     )

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

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

        snake_case__    							=				np.array([0.4992, 0.5101, 0.5004, 0.3949, 0.3604, 0.4735, 0.5216, 0.5204, 0.4913]     )

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

       def 		_SCREAMING_SNAKE_CASE(			self     ):






        '''simple docstring'''

        snake_case__    							=				"cpu"  # ensure determinism for the device-dependent torch.Generator
        snake_case__    							=				self.get_dummy_components()
        snake_case__    							=				EulerAncestralDiscreteScheduler(
            beta_start=0.00085   ,					beta_end=0.012   ,					beta_schedule="scaled_linear"     )
        snake_case__    							=				StableDiffusionModelEditingPipeline(**__SCREAMING_SNAKE_CASE     )
        snake_case__    							=				sd_pipe.to(__SCREAMING_SNAKE_CASE     )
        sd_pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE     )

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

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

        snake_case__    							=				np.array([0.4747, 0.5372, 0.4779, 0.4982, 0.5543, 0.4816, 0.5238, 0.4904, 0.5027]     )

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

       def 		_SCREAMING_SNAKE_CASE(			self     ):






        '''simple docstring'''

        snake_case__    							=				"cpu"  # ensure determinism for the device-dependent torch.Generator
        snake_case__    							=				self.get_dummy_components()
        snake_case__    							=				PNDMScheduler()
        snake_case__    							=				StableDiffusionModelEditingPipeline(**__SCREAMING_SNAKE_CASE     )
        snake_case__    							=				sd_pipe.to(__SCREAMING_SNAKE_CASE     )
        sd_pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE     )

        snake_case__    							=				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     ):
         snake_case__    							=				sd_pipe(**__SCREAMING_SNAKE_CASE     ).images

       def 		_SCREAMING_SNAKE_CASE(			self     ):






        '''simple docstring'''

        super().test_inference_batch_single_identical(expected_max_diff=5e-3     )

       def 		_SCREAMING_SNAKE_CASE(			self     ):






        '''simple docstring'''

        super().test_attention_slicing_forward_pass(expected_max_diff=5e-3     )



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

       def 		_SCREAMING_SNAKE_CASE(			self     ):






        '''simple docstring'''

        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()

       def 		_SCREAMING_SNAKE_CASE(			self   ,					__SCREAMING_SNAKE_CASE=0     ):






        '''simple docstring'''

        snake_case__    							=				torch.manual_seed(__SCREAMING_SNAKE_CASE     )
        snake_case__    							=				{
            "prompt": "A field of roses",
            "generator": generator,
            "num_inference_steps": 3,
            "guidance_scale": 7.5,
            "output_type": "numpy",
        }
        return inputs

       def 		_SCREAMING_SNAKE_CASE(			self     ):






        '''simple docstring'''

        snake_case__    							=				"CompVis/stable-diffusion-v1-4"
        snake_case__    							=				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()

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

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

        snake_case__    							=				np.array(
            [0.6749496, 0.6386453, 0.51443267, 0.66094905, 0.61921215, 0.5491332, 0.5744417, 0.58075106, 0.5174658]     )

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

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

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

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

       def 		_SCREAMING_SNAKE_CASE(			self     ):






        '''simple docstring'''

        torch.cuda.empty_cache()
        torch.cuda.reset_max_memory_allocated()
        torch.cuda.reset_peak_memory_stats()

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

        snake_case__    							=				self.get_inputs()
        snake_case__    							=				pipe(**__SCREAMING_SNAKE_CASE     )

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








"""simple docstring"""


import json
import os
import random
import shutil
import tempfile
import unittest

import numpy as np
import pytest

from transformers import BertTokenizer, BertTokenizerFast
from transformers.models.bert.tokenization_bert import VOCAB_FILES_NAMES
from transformers.testing_utils import require_vision
from transformers.utils import IMAGE_PROCESSOR_NAME, is_vision_available


if is_vision_available():
   from PIL import Image

   from transformers import FlavaImageProcessor, FlavaProcessor
   from transformers.models.flava.image_processing_flava import (
       FLAVA_CODEBOOK_MEAN,
       FLAVA_CODEBOOK_STD,
       FLAVA_IMAGE_MEAN,
       FLAVA_IMAGE_STD,
   )



@require_vision
class      lowercase_				(     unittest.TestCase							):

       def 		_SCREAMING_SNAKE_CASE(			self     ):






        '''simple docstring'''

        snake_case__    							=				tempfile.mkdtemp()

        # fmt: off
        snake_case__    							=				["[UNK]", "[CLS]", "[SEP]", "[PAD]", "[MASK]", "want", "##want", "##ed", "wa", "un", "runn", "##ing", ",", "low", "lowest"]
        # fmt: on
        snake_case__    							=				os.path.join(self.tmpdirname   ,					VOCAB_FILES_NAMES["vocab_file"]     )

        with open(self.vocab_file   ,					"w"   ,					encoding="utf-8"     ) as fp:
         fp.write("".join([x + "\n" for x in vocab_tokens]     )     )

        snake_case__    							=				{
            "image_mean": FLAVA_IMAGE_MEAN,
            "image_std": FLAVA_IMAGE_STD,
            "do_normalize": True,
            "do_resize": True,
            "size": 224,
            "do_center_crop": True,
            "crop_size": 224,
            "input_size_patches": 14,
            "total_mask_patches": 75,
            "mask_group_max_patches": None,
            "mask_group_min_patches": 16,
            "mask_group_min_aspect_ratio": 0.3,
            "mask_group_max_aspect_ratio": None,
            "codebook_do_resize": True,
            "codebook_size": 112,
            "codebook_do_center_crop": True,
            "codebook_crop_size": 112,
            "codebook_do_map_pixels": True,
            "codebook_do_normalize": True,
            "codebook_image_mean": FLAVA_CODEBOOK_MEAN,
            "codebook_image_std": FLAVA_CODEBOOK_STD,
        }

        snake_case__    							=				os.path.join(self.tmpdirname   ,					__SCREAMING_SNAKE_CASE     )
        with open(self.image_processor_file   ,					"w"   ,					encoding="utf-8"     ) as fp:
         json.dump(__SCREAMING_SNAKE_CASE   ,					__SCREAMING_SNAKE_CASE     )

       def 		_SCREAMING_SNAKE_CASE(			self   ,					**__SCREAMING_SNAKE_CASE     ):






        '''simple docstring'''

        return BertTokenizer.from_pretrained(self.tmpdirname   ,					**__SCREAMING_SNAKE_CASE     )

       def 		_SCREAMING_SNAKE_CASE(			self   ,					**__SCREAMING_SNAKE_CASE     ):






        '''simple docstring'''

        return BertTokenizerFast.from_pretrained(self.tmpdirname   ,					**__SCREAMING_SNAKE_CASE     )

       def 		_SCREAMING_SNAKE_CASE(			self   ,					**__SCREAMING_SNAKE_CASE     ):






        '''simple docstring'''

        return FlavaImageProcessor.from_pretrained(self.tmpdirname   ,					**__SCREAMING_SNAKE_CASE     )

       def 		_SCREAMING_SNAKE_CASE(			self     ):






        '''simple docstring'''

        shutil.rmtree(self.tmpdirname     )

       def 		_SCREAMING_SNAKE_CASE(			self     ):






        '''simple docstring'''

        snake_case__    							=				[np.random.randint(255   ,					size=(3, 30, 400)   ,					dtype=np.uinta     )]

        snake_case__    							=				[Image.fromarray(np.moveaxis(__SCREAMING_SNAKE_CASE   ,					0   ,					-1     )     ) for x in image_inputs]

        return image_inputs

       def 		_SCREAMING_SNAKE_CASE(			self     ):






        '''simple docstring'''

        snake_case__    							=				self.get_tokenizer()
        snake_case__    							=				self.get_rust_tokenizer()
        snake_case__    							=				self.get_image_processor()

        snake_case__    							=				FlavaProcessor(tokenizer=__SCREAMING_SNAKE_CASE   ,					image_processor=__SCREAMING_SNAKE_CASE     )
        processor_slow.save_pretrained(self.tmpdirname     )
        snake_case__    							=				FlavaProcessor.from_pretrained(self.tmpdirname   ,					use_fast=__SCREAMING_SNAKE_CASE     )

        snake_case__    							=				FlavaProcessor(tokenizer=__SCREAMING_SNAKE_CASE   ,					image_processor=__SCREAMING_SNAKE_CASE     )
        processor_fast.save_pretrained(self.tmpdirname     )
        snake_case__    							=				FlavaProcessor.from_pretrained(self.tmpdirname     )

        self.assertEqual(processor_slow.tokenizer.get_vocab()   ,					tokenizer_slow.get_vocab()     )
        self.assertEqual(processor_fast.tokenizer.get_vocab()   ,					tokenizer_fast.get_vocab()     )
        self.assertEqual(tokenizer_slow.get_vocab()   ,					tokenizer_fast.get_vocab()     )
        self.assertIsInstance(processor_slow.tokenizer   ,					__SCREAMING_SNAKE_CASE     )
        self.assertIsInstance(processor_fast.tokenizer   ,					__SCREAMING_SNAKE_CASE     )

        self.assertEqual(processor_slow.image_processor.to_json_string()   ,					image_processor.to_json_string()     )
        self.assertEqual(processor_fast.image_processor.to_json_string()   ,					image_processor.to_json_string()     )
        self.assertIsInstance(processor_slow.image_processor   ,					__SCREAMING_SNAKE_CASE     )
        self.assertIsInstance(processor_fast.image_processor   ,					__SCREAMING_SNAKE_CASE     )

       def 		_SCREAMING_SNAKE_CASE(			self     ):






        '''simple docstring'''

        snake_case__    							=				FlavaProcessor(tokenizer=self.get_tokenizer()   ,					image_processor=self.get_image_processor()     )
        processor.save_pretrained(self.tmpdirname     )

        snake_case__    							=				self.get_tokenizer(bos_token="(BOS)"   ,					eos_token="(EOS)"     )
        snake_case__    							=				self.get_image_processor(do_normalize=__SCREAMING_SNAKE_CASE   ,					padding_value=1.0     )

        snake_case__    							=				FlavaProcessor.from_pretrained(
            self.tmpdirname   ,					bos_token="(BOS)"   ,					eos_token="(EOS)"   ,					do_normalize=__SCREAMING_SNAKE_CASE   ,					padding_value=1.0     )

        self.assertEqual(processor.tokenizer.get_vocab()   ,					tokenizer_add_kwargs.get_vocab()     )
        self.assertIsInstance(processor.tokenizer   ,					__SCREAMING_SNAKE_CASE     )

        self.assertEqual(processor.image_processor.to_json_string()   ,					image_processor_add_kwargs.to_json_string()     )
        self.assertIsInstance(processor.image_processor   ,					__SCREAMING_SNAKE_CASE     )

       def 		_SCREAMING_SNAKE_CASE(			self     ):






        '''simple docstring'''

        snake_case__    							=				self.get_image_processor()
        snake_case__    							=				self.get_tokenizer()

        snake_case__    							=				FlavaProcessor(tokenizer=__SCREAMING_SNAKE_CASE   ,					image_processor=__SCREAMING_SNAKE_CASE     )

        snake_case__    							=				self.prepare_image_inputs()

        snake_case__    							=				image_processor(__SCREAMING_SNAKE_CASE   ,					return_tensors="np"     )
        snake_case__    							=				processor(images=__SCREAMING_SNAKE_CASE   ,					return_tensors="np"     )

        for key in input_feat_extract.keys():
         self.assertAlmostEqual(input_feat_extract[key].sum()   ,					input_processor[key].sum()   ,					delta=1e-2     )

        # With rest of the args
        random.seed(1234     )
        snake_case__    							=				image_processor(
            __SCREAMING_SNAKE_CASE   ,					return_image_mask=__SCREAMING_SNAKE_CASE   ,					return_codebook_pixels=__SCREAMING_SNAKE_CASE   ,					return_tensors="np"     )
        random.seed(1234     )
        snake_case__    							=				processor(
            images=__SCREAMING_SNAKE_CASE   ,					return_image_mask=__SCREAMING_SNAKE_CASE   ,					return_codebook_pixels=__SCREAMING_SNAKE_CASE   ,					return_tensors="np"     )

        for key in input_feat_extract.keys():
         self.assertAlmostEqual(input_feat_extract[key].sum()   ,					input_processor[key].sum()   ,					delta=1e-2     )

       def 		_SCREAMING_SNAKE_CASE(			self     ):






        '''simple docstring'''

        snake_case__    							=				self.get_image_processor()
        snake_case__    							=				self.get_tokenizer()

        snake_case__    							=				FlavaProcessor(tokenizer=__SCREAMING_SNAKE_CASE   ,					image_processor=__SCREAMING_SNAKE_CASE     )

        snake_case__    							=				"lower newer"

        snake_case__    							=				processor(text=__SCREAMING_SNAKE_CASE     )

        snake_case__    							=				tokenizer(__SCREAMING_SNAKE_CASE     )

        for key in encoded_tok.keys():
         self.assertListEqual(encoded_tok[key]   ,					encoded_processor[key]     )

       def 		_SCREAMING_SNAKE_CASE(			self     ):






        '''simple docstring'''

        snake_case__    							=				self.get_image_processor()
        snake_case__    							=				self.get_tokenizer()

        snake_case__    							=				FlavaProcessor(tokenizer=__SCREAMING_SNAKE_CASE   ,					image_processor=__SCREAMING_SNAKE_CASE     )

        snake_case__    							=				"lower newer"
        snake_case__    							=				self.prepare_image_inputs()

        snake_case__    							=				processor(text=__SCREAMING_SNAKE_CASE   ,					images=__SCREAMING_SNAKE_CASE     )

        self.assertListEqual(list(inputs.keys()     )   ,					["input_ids", "token_type_ids", "attention_mask", "pixel_values"]     )

        # add extra args
        snake_case__    							=				processor(text=__SCREAMING_SNAKE_CASE   ,					images=__SCREAMING_SNAKE_CASE   ,					return_codebook_pixels=__SCREAMING_SNAKE_CASE   ,					return_image_mask=__SCREAMING_SNAKE_CASE     )

        self.assertListEqual(
            list(inputs.keys()     )   ,					[
                "input_ids",
                "token_type_ids",
                "attention_mask",
                "pixel_values",
                "codebook_pixel_values",
                "bool_masked_pos",
            ]   ,					)

        # test if it raises when no input is passed
        with pytest.raises(__SCREAMING_SNAKE_CASE     ):
         processor()

       def 		_SCREAMING_SNAKE_CASE(			self     ):






        '''simple docstring'''

        snake_case__    							=				self.get_image_processor()
        snake_case__    							=				self.get_tokenizer()

        snake_case__    							=				FlavaProcessor(tokenizer=__SCREAMING_SNAKE_CASE   ,					image_processor=__SCREAMING_SNAKE_CASE     )

        snake_case__    							=				[[1, 4, 5, 8, 1, 0, 8], [3, 4, 3, 1, 1, 8, 9]]

        snake_case__    							=				processor.batch_decode(__SCREAMING_SNAKE_CASE     )
        snake_case__    							=				tokenizer.batch_decode(__SCREAMING_SNAKE_CASE     )

        self.assertListEqual(__SCREAMING_SNAKE_CASE   ,					__SCREAMING_SNAKE_CASE     )

       def 		_SCREAMING_SNAKE_CASE(			self     ):






        '''simple docstring'''

        snake_case__    							=				self.get_image_processor()
        snake_case__    							=				self.get_tokenizer()

        snake_case__    							=				FlavaProcessor(tokenizer=__SCREAMING_SNAKE_CASE   ,					image_processor=__SCREAMING_SNAKE_CASE     )

        snake_case__    							=				"lower newer"
        snake_case__    							=				self.prepare_image_inputs()

        snake_case__    							=				processor(text=__SCREAMING_SNAKE_CASE   ,					images=__SCREAMING_SNAKE_CASE     )

        self.assertListEqual(list(inputs.keys()     )   ,					processor.model_input_names     )

