

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 					__lowerCAmelCase	(   SCREAMING_SNAKE_CASE     ,     SCREAMING_SNAKE_CASE     ,     unittest.TestCase				):
  _a		     =    StableDiffusionDiffEditPipeline
  _a		     =    TEXT_GUIDED_IMAGE_INPAINTING_PARAMS - {"""height""", """width""", """image"""} | {"""image_latents"""}
  _a		     =    TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS - {"""image"""} | {"""image_latents"""}
  _a		     =    frozenset(
      []				)  # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
  _a		     =    frozenset([]				)




  def     A__	(							self   )    ->					List[Any]:
        '''simple docstring'''






        torch.manual_seed(0   )
        _lowercase      			=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=lowerCAmelCase				,  )
        _lowercase      			=DDIMScheduler(
            beta_start=0.00085				,  beta_end=0.012				,  beta_schedule='scaled_linear'				,  clip_sample=lowerCAmelCase				,  set_alpha_to_one=lowerCAmelCase				,  )
        _lowercase      			=DDIMInverseScheduler(
            beta_start=0.00085				,  beta_end=0.012				,  beta_schedule='scaled_linear'				,  clip_sample=lowerCAmelCase				,  set_alpha_to_zero=lowerCAmelCase				,  )
        torch.manual_seed(0   )
        _lowercase      			=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   )
        _lowercase      			=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=1_000				,  hidden_act='gelu'				,  projection_dim=512				,  )
        _lowercase      			=CLIPTextModel(lowerCAmelCase   )
        _lowercase      			=CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'   )

        _lowercase      			={
            '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     A__	(							self				,  lowerCAmelCase				,  lowerCAmelCase=0   )    ->					Tuple:
        '''simple docstring'''






        _lowercase      			=floats_tensor((1, 16, 16)				,  rng=random.Random(lowerCAmelCase   )   ).to(lowerCAmelCase   )
        _lowercase      			=floats_tensor((1, 2, 4, 16, 16)				,  rng=random.Random(lowerCAmelCase   )   ).to(lowerCAmelCase   )
        if str(lowerCAmelCase   ).startswith('mps'   ):
              _lowercase      			=torch.manual_seed(lowerCAmelCase   )
        else:
              _lowercase      			=torch.Generator(device=lowerCAmelCase   ).manual_seed(lowerCAmelCase   )
        _lowercase      			={
            '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     A__	(							self				,  lowerCAmelCase				,  lowerCAmelCase=0   )    ->					Optional[Any]:
        '''simple docstring'''






        _lowercase      			=floats_tensor((1, 3, 32, 32)				,  rng=random.Random(lowerCAmelCase   )   ).to(lowerCAmelCase   )
        _lowercase      			=image.cpu().permute(0				,  2				,  3				,  1   )[0]
        _lowercase      			=Image.fromarray(np.uinta(lowerCAmelCase   )   ).convert('RGB'   )
        if str(lowerCAmelCase   ).startswith('mps'   ):
              _lowercase      			=torch.manual_seed(lowerCAmelCase   )
        else:
              _lowercase      			=torch.Generator(device=lowerCAmelCase   ).manual_seed(lowerCAmelCase   )
        _lowercase      			={
            '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     A__	(							self				,  lowerCAmelCase				,  lowerCAmelCase=0   )    ->					List[Any]:
        '''simple docstring'''






        _lowercase      			=floats_tensor((1, 3, 32, 32)				,  rng=random.Random(lowerCAmelCase   )   ).to(lowerCAmelCase   )
        _lowercase      			=image.cpu().permute(0				,  2				,  3				,  1   )[0]
        _lowercase      			=Image.fromarray(np.uinta(lowerCAmelCase   )   ).convert('RGB'   )
        if str(lowerCAmelCase   ).startswith('mps'   ):
              _lowercase      			=torch.manual_seed(lowerCAmelCase   )
        else:
              _lowercase      			=torch.Generator(device=lowerCAmelCase   ).manual_seed(lowerCAmelCase   )
        _lowercase      			={
            '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     A__	(							self   )    ->					Any:
        '''simple docstring'''






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

        _lowercase      			=self.get_dummy_components()
        _lowercase      			=self.pipeline_class(**lowerCAmelCase   )
        pipe.to(lowerCAmelCase   )
        pipe.set_progress_bar_config(disable=lowerCAmelCase   )

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

        _lowercase      			=self.get_dummy_inputs(lowerCAmelCase   )
        _lowercase      			=pipe(**lowerCAmelCase   )[0]

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

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

        _lowercase      			=self.get_dummy_inputs(lowerCAmelCase   )
        _lowercase      			=pipe_loaded(**lowerCAmelCase   )[0]

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




  def     A__	(							self   )    ->					int:
        '''simple docstring'''






        _lowercase      			='cpu'

        _lowercase      			=self.get_dummy_components()
        _lowercase      			=self.pipeline_class(**lowerCAmelCase   )
        pipe.to(lowerCAmelCase   )
        pipe.set_progress_bar_config(disable=lowerCAmelCase   )

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

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




  def     A__	(							self   )    ->					Dict:
        '''simple docstring'''






        _lowercase      			='cpu'

        _lowercase      			=self.get_dummy_components()
        _lowercase      			=self.pipeline_class(**lowerCAmelCase   )
        pipe.to(lowerCAmelCase   )
        pipe.set_progress_bar_config(disable=lowerCAmelCase   )

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

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




  def     A__	(							self   )    ->					List[Any]:
        '''simple docstring'''






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




  def     A__	(							self   )    ->					Any:
        '''simple docstring'''






        _lowercase      			='cpu'

        _lowercase      			=self.get_dummy_components()

        _lowercase      			={'beta_start': 0.00085, 'beta_end': 0.012, 'beta_schedule': 'scaled_linear'}
        _lowercase      			=DPMSolverMultistepScheduler(**lowerCAmelCase   )
        _lowercase      			=DPMSolverMultistepInverseScheduler(**lowerCAmelCase   )

        _lowercase      			=self.pipeline_class(**lowerCAmelCase   )
        pipe.to(lowerCAmelCase   )
        pipe.set_progress_bar_config(disable=lowerCAmelCase   )

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

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





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




  def     A__	(							self   )    ->					Union[str, Any]:
        '''simple docstring'''






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




  @classmethod
  def     A__	(							cls   )    ->					int:
        '''simple docstring'''






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

        _lowercase      			=raw_image.convert('RGB'   ).resize((768, 768)   )

        _lowercase      			=raw_image




  def     A__	(							self   )    ->					Dict:
        '''simple docstring'''






        _lowercase      			=torch.manual_seed(0   )

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

        _lowercase      			='a bowl of fruit'
        _lowercase      			='a bowl of pears'

        _lowercase      			=pipe.generate_mask(
            image=self.raw_image				,  source_prompt=lowerCAmelCase				,  target_prompt=lowerCAmelCase				,  generator=lowerCAmelCase				,  )

        _lowercase      			=pipe.invert(
            prompt=lowerCAmelCase				,  image=self.raw_image				,  inpaint_strength=0.7				,  generator=lowerCAmelCase   ).latents

        _lowercase      			=pipe(
            prompt=lowerCAmelCase				,  mask_image=lowerCAmelCase				,  image_latents=lowerCAmelCase				,  generator=lowerCAmelCase				,  negative_prompt=lowerCAmelCase				,  inpaint_strength=0.7				,  output_type='numpy'				,  ).images[0]

        _lowercase      			=(
            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     A__	(							self   )    ->					Any:
        '''simple docstring'''






        _lowercase      			=torch.manual_seed(0   )

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

        _lowercase      			='a bowl of fruit'
        _lowercase      			='a bowl of pears'

        _lowercase      			=pipe.generate_mask(
            image=self.raw_image				,  source_prompt=lowerCAmelCase				,  target_prompt=lowerCAmelCase				,  generator=lowerCAmelCase				,  )

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

        _lowercase      			=pipe(
            prompt=lowerCAmelCase				,  mask_image=lowerCAmelCase				,  image_latents=lowerCAmelCase				,  generator=lowerCAmelCase				,  negative_prompt=lowerCAmelCase				,  inpaint_strength=0.7				,  num_inference_steps=25				,  output_type='numpy'				,  ).images[0]

        _lowercase      			=(
            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





import json
import os
import tempfile

import transformers

import datasets
from utils import generate_example_dataset, get_duration


lowercase_   							=  5_0_0_0_0_0

lowercase_							,		lowercase_   							=  os.path.split(__file__)
lowercase_   							=  os.path.join(RESULTS_BASEPATH, 'results', RESULTS_FILENAME.replace('.py', '.json'))
@get_duration
def 	a (     A__      :  datasets.Dataset ,  **A__      :  List[str]	)       -> Optional[Any]:





      """simple docstring"""



      _lowercase      			=dataset.map(**A__	)
@get_duration
def 	a (     A__      :  datasets.Dataset ,  **A__      :  List[Any]	)       -> List[str]:





      """simple docstring"""



      _lowercase      			=dataset.filter(**A__	)
def 	a (     )       -> Union[str, Any]:





      """simple docstring"""



      _lowercase      			={'num examples': SPEED_TEST_N_EXAMPLES}
      with tempfile.TemporaryDirectory() as tmp_dir:
            _lowercase      			=datasets.Features({'text': datasets.Value('string'	), 'numbers': datasets.Value('float32'	)}	)
            _lowercase      			=generate_example_dataset(
                os.path.join(A__ ,  'dataset.arrow'	) ,  A__ ,  num_examples=A__	)

            _lowercase      			=transformers.AutoTokenizer.from_pretrained('bert-base-cased' ,  use_fast=A__	)

            def tokenize(A__      :  int	):
                  return tokenizer(examples['text']	)

            _lowercase      			=map(A__	)

            _lowercase      			=map(A__ ,  batched=A__	)

            _lowercase      			=map(A__ ,  function=lambda A__	: None ,  batched=A__	)

            with dataset.formatted_as(type='numpy'	):
                  _lowercase      			=map(A__ ,  function=lambda A__	: None ,  batched=A__	)

            with dataset.formatted_as(type='pandas'	):
                  _lowercase      			=map(A__ ,  function=lambda A__	: None ,  batched=A__	)

            with dataset.formatted_as(type='torch' ,  columns='numbers'	):
                  _lowercase      			=map(A__ ,  function=lambda A__	: None ,  batched=A__	)

            with dataset.formatted_as(type='tensorflow' ,  columns='numbers'	):
                  _lowercase      			=map(A__ ,  function=lambda A__	: None ,  batched=A__	)

            _lowercase      			=map(A__ ,  function=A__ ,  batched=A__	)

            _lowercase      			=filter(A__	)

            # Activate later when tokenizer support batched inputs
            # with dataset.formatted_as(type='numpy'):
            #     times[func.__name__ + " fast-tokenizer batched numpy"] = func(dataset, function=tokenize, batched=True)

      with open(A__ ,  'wb'	) as f:
            f.write(json.dumps(A__	).encode('utf-8'	)	)


if __name__ == "__main__":  # useful to run the profiler
  benchmark_map_filter()


