

'''simple docstring'''







import gc
import random
import tempfile
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMInverseScheduler,
    DDIMScheduler,
    DPMSolverMultistepInverseScheduler,
    DPMSolverMultistepScheduler,
    StableDiffusionDiffEditPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import load_image, slow
from diffusers.utils.testing_utils import enable_full_determinism, floats_tensor, require_torch_gpu, torch_device

from ..pipeline_params import TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS, TEXT_GUIDED_IMAGE_INPAINTING_PARAMS
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()


class        A__   (   _snake_case    ,    _snake_case    ,    unittest.TestCase      ):
    lowercase						       =		StableDiffusionDiffEditPipeline
    lowercase						       =		TEXT_GUIDED_IMAGE_INPAINTING_PARAMS - {"height", "width", "image"} | {"image_latents"}
    lowercase						       =		TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS - {"image"} | {"image_latents"}
    lowercase						       =		frozenset(
        []      )  # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
    lowercase						       =		frozenset([]      )



    def        snake_case_      (       self							)    ->   Dict:





           '''simple docstring'''
           torch.manual_seed(0							)
           A_						       =   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		,					attention_head_dim=(2, 4)		,					use_linear_projection=UpperCamelCase__		,					)
           A_						       =   DDIMScheduler(
               beta_start=0.00085		,					beta_end=0.012		,					beta_schedule="""scaled_linear"""		,					clip_sample=UpperCamelCase__		,					set_alpha_to_one=UpperCamelCase__		,					)
           A_						       =   DDIMInverseScheduler(
               beta_start=0.00085		,					beta_end=0.012		,					beta_schedule="""scaled_linear"""		,					clip_sample=UpperCamelCase__		,					set_alpha_to_zero=UpperCamelCase__		,					)
           torch.manual_seed(0							)
           A_						       =   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		,					sample_size=128		,					)
           torch.manual_seed(0							)
           A_						       =   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=1000		,					hidden_act="""gelu"""		,					projection_dim=512		,					)
           A_						       =   CLIPTextModel(UpperCamelCase__							)
           A_						       =   CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""							)

           A_						       =   {
               """unet""": unet,
               """scheduler""": scheduler,
               """inverse_scheduler""": inverse_scheduler,
               """vae""": vae,
               """text_encoder""": text_encoder,
               """tokenizer""": tokenizer,
               """safety_checker""": None,
               """feature_extractor""": None,
           }

           return components



    def        snake_case_      (       self		,					UpperCamelCase__		,					UpperCamelCase__=0							)    ->   Optional[Any]:





           '''simple docstring'''
           A_						       =   floats_tensor((1, 16, 16)		,					rng=random.Random(UpperCamelCase__							)							).to(UpperCamelCase__							)
           A_						       =   floats_tensor((1, 2, 4, 16, 16)		,					rng=random.Random(UpperCamelCase__							)							).to(UpperCamelCase__							)
           if str(UpperCamelCase__							).startswith("""mps"""							):
                  A_						       =   torch.manual_seed(UpperCamelCase__							)
           else:
                  A_						       =   torch.Generator(device=UpperCamelCase__							).manual_seed(UpperCamelCase__							)
           A_						       =   {
               """prompt""": """a dog and a newt""",
               """mask_image""": mask,
               """image_latents""": latents,
               """generator""": generator,
               """num_inference_steps""": 2,
               """inpaint_strength""": 1.0,
               """guidance_scale""": 6.0,
               """output_type""": """numpy""",
           }

           return inputs



    def        snake_case_      (       self		,					UpperCamelCase__		,					UpperCamelCase__=0							)    ->   str:





           '''simple docstring'''
           A_						       =   floats_tensor((1, 3, 32, 32)		,					rng=random.Random(UpperCamelCase__							)							).to(UpperCamelCase__							)
           A_						       =   image.cpu().permute(0		,					2		,					3		,					1							)[0]
           A_						       =   Image.fromarray(np.uinta(UpperCamelCase__							)							).convert("""RGB"""							)
           if str(UpperCamelCase__							).startswith("""mps"""							):
                  A_						       =   torch.manual_seed(UpperCamelCase__							)
           else:
                  A_						       =   torch.Generator(device=UpperCamelCase__							).manual_seed(UpperCamelCase__							)
           A_						       =   {
               """image""": image,
               """source_prompt""": """a cat and a frog""",
               """target_prompt""": """a dog and a newt""",
               """generator""": generator,
               """num_inference_steps""": 2,
               """num_maps_per_mask""": 2,
               """mask_encode_strength""": 1.0,
               """guidance_scale""": 6.0,
               """output_type""": """numpy""",
           }

           return inputs



    def        snake_case_      (       self		,					UpperCamelCase__		,					UpperCamelCase__=0							)    ->   Tuple:





           '''simple docstring'''
           A_						       =   floats_tensor((1, 3, 32, 32)		,					rng=random.Random(UpperCamelCase__							)							).to(UpperCamelCase__							)
           A_						       =   image.cpu().permute(0		,					2		,					3		,					1							)[0]
           A_						       =   Image.fromarray(np.uinta(UpperCamelCase__							)							).convert("""RGB"""							)
           if str(UpperCamelCase__							).startswith("""mps"""							):
                  A_						       =   torch.manual_seed(UpperCamelCase__							)
           else:
                  A_						       =   torch.Generator(device=UpperCamelCase__							).manual_seed(UpperCamelCase__							)
           A_						       =   {
               """image""": image,
               """prompt""": """a cat and a frog""",
               """generator""": generator,
               """num_inference_steps""": 2,
               """inpaint_strength""": 1.0,
               """guidance_scale""": 6.0,
               """decode_latents""": True,
               """output_type""": """numpy""",
           }
           return inputs



    def        snake_case_      (       self							)    ->   str:





           '''simple docstring'''
           if not hasattr(self.pipeline_class		,					"""_optional_components"""							):
                  return

           A_						       =   self.get_dummy_components()
           A_						       =   self.pipeline_class(**UpperCamelCase__							)
           pipe.to(UpperCamelCase__							)
           pipe.set_progress_bar_config(disable=UpperCamelCase__							)

           # set all optional components to None and update pipeline config accordingly
           for optional_component in pipe._optional_components:
                  setattr(UpperCamelCase__		,					UpperCamelCase__		,					UpperCamelCase__							)
           pipe.register_modules(**{optional_component: None for optional_component in pipe._optional_components}							)

           A_						       =   self.get_dummy_inputs(UpperCamelCase__							)
           A_						       =   pipe(**UpperCamelCase__							)[0]

           with tempfile.TemporaryDirectory() as tmpdir:
                  pipe.save_pretrained(UpperCamelCase__							)
                  A_						       =   self.pipeline_class.from_pretrained(UpperCamelCase__							)
                  pipe_loaded.to(UpperCamelCase__							)
                  pipe_loaded.set_progress_bar_config(disable=UpperCamelCase__							)

           for optional_component in pipe._optional_components:
                  self.assertTrue(
                      getattr(UpperCamelCase__		,					UpperCamelCase__							) is None		,					f'''`{optional_component}` did not stay set to None after loading.'''		,					)

           A_						       =   self.get_dummy_inputs(UpperCamelCase__							)
           A_						       =   pipe_loaded(**UpperCamelCase__							)[0]

           A_						       =   np.abs(output - output_loaded							).max()
           self.assertLess(UpperCamelCase__		,					1e-4							)



    def        snake_case_      (       self							)    ->   Tuple:





           '''simple docstring'''
           A_						       =   """cpu"""

           A_						       =   self.get_dummy_components()
           A_						       =   self.pipeline_class(**UpperCamelCase__							)
           pipe.to(UpperCamelCase__							)
           pipe.set_progress_bar_config(disable=UpperCamelCase__							)

           A_						       =   self.get_dummy_mask_inputs(UpperCamelCase__							)
           A_						       =   pipe.generate_mask(**UpperCamelCase__							)
           A_						       =   mask[0, -3:, -3:]

           self.assertEqual(mask.shape		,					(1, 16, 16)							)
           A_						       =   np.array([0] * 9							)
           A_						       =   np.abs(mask_slice.flatten() - expected_slice							).max()
           self.assertLessEqual(UpperCamelCase__		,					1e-3							)
           self.assertEqual(mask[0, -3, -4]		,					0							)



    def        snake_case_      (       self							)    ->   int:





           '''simple docstring'''
           A_						       =   """cpu"""

           A_						       =   self.get_dummy_components()
           A_						       =   self.pipeline_class(**UpperCamelCase__							)
           pipe.to(UpperCamelCase__							)
           pipe.set_progress_bar_config(disable=UpperCamelCase__							)

           A_						       =   self.get_dummy_inversion_inputs(UpperCamelCase__							)
           A_						       =   pipe.invert(**UpperCamelCase__							).images
           A_						       =   image[0, -1, -3:, -3:]

           self.assertEqual(image.shape		,					(2, 32, 32, 3)							)
           A_						       =   np.array(
               [0.5150, 0.5134, 0.5043, 0.5376, 0.4694, 0.51050, 0.5015, 0.4407, 0.4799]		,					)
           A_						       =   np.abs(image_slice.flatten() - expected_slice							).max()
           self.assertLessEqual(UpperCamelCase__		,					1e-3							)



    def        snake_case_      (       self							)    ->   int:





           '''simple docstring'''
           super().test_inference_batch_single_identical(expected_max_diff=5e-3							)







    def        snake_case_      (       self							)    ->   int:





           '''simple docstring'''
           A_						       =   """cpu"""

           A_						       =   self.get_dummy_components()

           A_						       =   {"""beta_start""": 0.00085, """beta_end""": 0.012, """beta_schedule""": """scaled_linear"""}
           A_						       =   DPMSolverMultistepScheduler(**UpperCamelCase__							)
           A_						       =   DPMSolverMultistepInverseScheduler(**UpperCamelCase__							)

           A_						       =   self.pipeline_class(**UpperCamelCase__							)
           pipe.to(UpperCamelCase__							)
           pipe.set_progress_bar_config(disable=UpperCamelCase__							)

           A_						       =   self.get_dummy_inversion_inputs(UpperCamelCase__							)
           A_						       =   pipe.invert(**UpperCamelCase__							).images
           A_						       =   image[0, -1, -3:, -3:]

           self.assertEqual(image.shape		,					(2, 32, 32, 3)							)
           A_						       =   np.array(
               [0.5150, 0.5134, 0.5043, 0.5376, 0.4694, 0.51050, 0.5015, 0.4407, 0.4799]		,					)
           A_						       =   np.abs(image_slice.flatten() - expected_slice							).max()
           self.assertLessEqual(UpperCamelCase__		,					1e-3							)


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



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





           '''simple docstring'''
           super().tearDown()
           gc.collect()
           torch.cuda.empty_cache()



    @classmethod
    def        snake_case_      (       cls							)    ->   int:





           '''simple docstring'''
           A_						       =   load_image(
               """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/diffedit/fruit.png"""							)

           A_						       =   raw_image.convert("""RGB"""							).resize((768, 768)							)

           A_						       =   raw_image



    def        snake_case_      (       self							)    ->   Any:





           '''simple docstring'''
           A_						       =   torch.manual_seed(0							)

           A_						       =   StableDiffusionDiffEditPipeline.from_pretrained(
               """stabilityai/stable-diffusion-2-1"""		,					safety_checker=UpperCamelCase__		,					torch_dtype=torch.floataa							)
           A_						       =   DDIMScheduler.from_config(pipe.scheduler.config							)
           A_						       =   DDIMInverseScheduler.from_config(pipe.scheduler.config							)
           pipe.enable_model_cpu_offload()
           pipe.set_progress_bar_config(disable=UpperCamelCase__							)

           A_						       =   """a bowl of fruit"""
           A_						       =   """a bowl of pears"""

           A_						       =   pipe.generate_mask(
               image=self.raw_image		,					source_prompt=UpperCamelCase__		,					target_prompt=UpperCamelCase__		,					generator=UpperCamelCase__		,					)

           A_						       =   pipe.invert(
               prompt=UpperCamelCase__		,					image=self.raw_image		,					inpaint_strength=0.7		,					generator=UpperCamelCase__							).latents

           A_						       =   pipe(
               prompt=UpperCamelCase__		,					mask_image=UpperCamelCase__		,					image_latents=UpperCamelCase__		,					generator=UpperCamelCase__		,					negative_prompt=UpperCamelCase__		,					inpaint_strength=0.7		,					output_type="""numpy"""		,					).images[0]

           A_						       =   (
               np.array(
                   load_image(
                       """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
                       """/diffedit/pears.png"""							).resize((768, 768)							)							)
               / 255
           )
           assert np.abs((expected_image - image).max()							) < 5e-1







    def        snake_case_      (       self							)    ->   str:





           '''simple docstring'''
           A_						       =   torch.manual_seed(0							)

           A_						       =   StableDiffusionDiffEditPipeline.from_pretrained(
               """stabilityai/stable-diffusion-2-1"""		,					safety_checker=UpperCamelCase__		,					torch_dtype=torch.floataa							)
           A_						       =   DPMSolverMultistepScheduler.from_config(pipe.scheduler.config							)
           A_						       =   DPMSolverMultistepInverseScheduler.from_config(pipe.scheduler.config							)
           pipe.enable_model_cpu_offload()
           pipe.set_progress_bar_config(disable=UpperCamelCase__							)

           A_						       =   """a bowl of fruit"""
           A_						       =   """a bowl of pears"""

           A_						       =   pipe.generate_mask(
               image=self.raw_image		,					source_prompt=UpperCamelCase__		,					target_prompt=UpperCamelCase__		,					generator=UpperCamelCase__		,					)

           A_						       =   pipe.invert(
               prompt=UpperCamelCase__		,					image=self.raw_image		,					inpaint_strength=0.7		,					generator=UpperCamelCase__		,					num_inference_steps=25		,					).latents

           A_						       =   pipe(
               prompt=UpperCamelCase__		,					mask_image=UpperCamelCase__		,					image_latents=UpperCamelCase__		,					generator=UpperCamelCase__		,					negative_prompt=UpperCamelCase__		,					inpaint_strength=0.7		,					num_inference_steps=25		,					output_type="""numpy"""		,					).images[0]

           A_						       =   (
               np.array(
                   load_image(
                       """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
                       """/diffedit/pears.png"""							).resize((768, 768)							)							)
               / 255
           )
           assert np.abs((expected_image - image).max()							) < 5e-1



'''simple docstring'''







from __future__ import annotations

import unittest

import numpy as np

from transformers import BlipTextConfig
from transformers.testing_utils import require_tf, slow
from transformers.utils import is_tf_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask


if is_tf_available():
      import tensorflow as tf

      from transformers import TFBlipTextModel
      from transformers.models.blip.modeling_tf_blip import TF_BLIP_PRETRAINED_MODEL_ARCHIVE_LIST


class        A__   :



    def __init__(       self		,					UpperCamelCase__		,					UpperCamelCase__=12		,					UpperCamelCase__=7		,					UpperCamelCase__=True		,					UpperCamelCase__=True		,					UpperCamelCase__=True		,					UpperCamelCase__=99		,					UpperCamelCase__=32		,					UpperCamelCase__=32		,					UpperCamelCase__=2		,					UpperCamelCase__=4		,					UpperCamelCase__=37		,					UpperCamelCase__=0.1		,					UpperCamelCase__=0.1		,					UpperCamelCase__=512		,					UpperCamelCase__=0.02		,					UpperCamelCase__=0		,					UpperCamelCase__=None		,					)    ->   Union[str, Any]:





           '''simple docstring'''
           A_						       =   parent
           A_						       =   batch_size
           A_						       =   seq_length
           A_						       =   is_training
           A_						       =   use_input_mask
           A_						       =   use_labels
           A_						       =   vocab_size
           A_						       =   hidden_size
           A_						       =   projection_dim
           A_						       =   num_hidden_layers
           A_						       =   num_attention_heads
           A_						       =   intermediate_size
           A_						       =   dropout
           A_						       =   attention_dropout
           A_						       =   max_position_embeddings
           A_						       =   initializer_range
           A_						       =   scope
           A_						       =   bos_token_id



    def        snake_case_      (       self							)    ->   Dict:





           '''simple docstring'''
           A_						       =   ids_tensor([self.batch_size, self.seq_length]		,					self.vocab_size							)

           A_						       =   None
           if self.use_input_mask:
                  A_						       =   random_attention_mask([self.batch_size, self.seq_length]							)

           if input_mask is not None:
                  A_						       =   input_mask.numpy()
                  A_       ,      A_						       =   input_mask.shape
                  A_						       =   np.random.randint(1		,					seq_length - 1		,					size=(batch_size,)							)
                  for batch_idx, start_index in enumerate(UpperCamelCase__							):
                         A_						       =   1
                         A_						       =   0

           A_						       =   self.get_config()

           return config, input_ids, tf.convert_to_tensor(UpperCamelCase__							)



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





           '''simple docstring'''
           return BlipTextConfig(
               vocab_size=self.vocab_size		,					hidden_size=self.hidden_size		,					projection_dim=self.projection_dim		,					num_hidden_layers=self.num_hidden_layers		,					num_attention_heads=self.num_attention_heads		,					intermediate_size=self.intermediate_size		,					dropout=self.dropout		,					attention_dropout=self.attention_dropout		,					max_position_embeddings=self.max_position_embeddings		,					initializer_range=self.initializer_range		,					bos_token_id=self.bos_token_id		,					)



    def        snake_case_      (       self		,					UpperCamelCase__		,					UpperCamelCase__		,					UpperCamelCase__							)    ->   int:





           '''simple docstring'''
           A_						       =   TFBlipTextModel(config=UpperCamelCase__							)
           A_						       =   model(UpperCamelCase__		,					attention_mask=UpperCamelCase__		,					training=UpperCamelCase__							)
           A_						       =   model(UpperCamelCase__		,					training=UpperCamelCase__							)
           self.parent.assertEqual(result.last_hidden_state.shape		,					(self.batch_size, self.seq_length, self.hidden_size)							)
           self.parent.assertEqual(result.pooler_output.shape		,					(self.batch_size, self.hidden_size)							)







    def        snake_case_      (       self							)    ->   List[Any]:





           '''simple docstring'''
           A_						       =   self.prepare_config_and_inputs()
           A_       ,      A_       ,      A_						       =   config_and_inputs
           A_						       =   {"""input_ids""": input_ids, """attention_mask""": input_mask}
           return config, inputs_dict


@require_tf
class        A__   (   _snake_case    ,    unittest.TestCase      ):
    lowercase						       =		(TFBlipTextModel,) if is_tf_available() else ()
    lowercase						       =		False
    lowercase						       =		False
    lowercase						       =		False



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





           '''simple docstring'''
           A_						       =   BlipTextModelTester(self							)
           A_						       =   ConfigTester(self		,					config_class=UpperCamelCase__		,					hidden_size=37							)



    def        snake_case_      (       self							)    ->   Tuple:





           '''simple docstring'''
           self.config_tester.run_common_tests()



    def        snake_case_      (       self							)    ->   str:





           '''simple docstring'''
           A_						       =   self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_model(*UpperCamelCase__							)



    def        snake_case_      (       self							)    ->   List[Any]:





           '''simple docstring'''
           pass



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





           '''simple docstring'''
           pass



    @unittest.skip(reason="""Blip does not use inputs_embeds"""							)
    def        snake_case_      (       self							)    ->   Union[str, Any]:





           '''simple docstring'''
           pass



    @unittest.skip(reason="""BlipTextModel has no base class and is not available in MODEL_MAPPING"""							)
    def        snake_case_      (       self							)    ->   List[Any]:





           '''simple docstring'''
           pass



    @unittest.skip(reason="""BlipTextModel has no base class and is not available in MODEL_MAPPING"""							)
    def        snake_case_      (       self							)    ->   Any:





           '''simple docstring'''
           pass



    @slow
    def        snake_case_      (       self							)    ->   List[Any]:





           '''simple docstring'''
           for model_name in TF_BLIP_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                  A_						       =   TFBlipTextModel.from_pretrained(UpperCamelCase__							)
                  self.assertIsNotNone(UpperCamelCase__							)







    def        snake_case_      (       self		,					UpperCamelCase__=True							)    ->   str:





           '''simple docstring'''
           super().test_pt_tf_model_equivalence(allow_missing_keys=UpperCamelCase__							)
