


# coding=utf-8
# Copyright 2023 The HuggingFace Inc. team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# this script dumps information about the environment

import os
import platform
import sys


SCREAMING_SNAKE_CASE		:Union[str, Any]     						=   '''3'''

print('''Python version:''', sys.version)

print('''OS platform:''', platform.platform())
print('''OS architecture:''', platform.machine())

try:
   import torch

   print('''Torch version:''', torch.__version__)
   print('''Cuda available:''', torch.cuda.is_available())
   print('''Cuda version:''', torch.version.cuda)
   print('''CuDNN version:''', torch.backends.cudnn.version())
   print('''Number of GPUs available:''', torch.cuda.device_count())
except ImportError:
   print('''Torch version:''', None)

try:
   import transformers

   print('''transformers version:''', transformers.__version__)
except ImportError:
   print('''transformers version:''', None)





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 		_lowercase						(					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 						SCREAMING_SNAKE_CASE__   ( self	:							List[Any]			)   ->						Dict:

  """simple docstring"""





  torch.manual_seed(0			)
  UpperCamelCase_					:						Dict =					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   ,      attention_head_dim=(2, 4)   ,      use_linear_projection=snake_case   ,      )
  UpperCamelCase_					:						Optional[int] =					DDIMScheduler(
      beta_start=0.00085   ,      beta_end=0.012   ,      beta_schedule='scaled_linear'   ,      clip_sample=snake_case   ,      set_alpha_to_one=snake_case   ,      )
  UpperCamelCase_					:						Tuple =					DDIMInverseScheduler(
      beta_start=0.00085   ,      beta_end=0.012   ,      beta_schedule='scaled_linear'   ,      clip_sample=snake_case   ,      set_alpha_to_zero=snake_case   ,      )
  torch.manual_seed(0			)
  UpperCamelCase_					:						List[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   ,      sample_size=1_2_8   ,      )
  torch.manual_seed(0			)
  UpperCamelCase_					:						List[Any] =					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   ,      hidden_act='gelu'   ,      projection_dim=5_1_2   ,      )
  UpperCamelCase_					:						List[str] =					CLIPTextModel(snake_case			)
  UpperCamelCase_					:						List[str] =					CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'			)

  UpperCamelCase_					:						Any =					{
      '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 						SCREAMING_SNAKE_CASE__   ( self	:							Optional[Any]   ,      snake_case	:							Dict   ,      snake_case	:							Tuple=0			)   ->						int:

  """simple docstring"""





  UpperCamelCase_					:						Optional[Any] =					floats_tensor((1, 1_6, 1_6)   ,      rng=random.Random(snake_case			)			).to(snake_case			)
  UpperCamelCase_					:						Optional[int] =					floats_tensor((1, 2, 4, 1_6, 1_6)   ,      rng=random.Random(snake_case			)			).to(snake_case			)
  if str(snake_case			).startswith('mps'			):
   UpperCamelCase_					:						int =					torch.manual_seed(snake_case			)
  else:
   UpperCamelCase_					:						Any =					torch.Generator(device=snake_case			).manual_seed(snake_case			)
  UpperCamelCase_					:						Union[str, Any] =					{
      '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 						SCREAMING_SNAKE_CASE__   ( self	:							Union[str, Any]   ,      snake_case	:							Union[str, Any]   ,      snake_case	:							Dict=0			)   ->						str:

  """simple docstring"""





  UpperCamelCase_					:						List[Any] =					floats_tensor((1, 3, 3_2, 3_2)   ,      rng=random.Random(snake_case			)			).to(snake_case			)
  UpperCamelCase_					:						int =					image.cpu().permute(0   ,      2   ,      3   ,      1			)[0]
  UpperCamelCase_					:						Optional[int] =					Image.fromarray(np.uinta(snake_case			)			).convert('RGB'			)
  if str(snake_case			).startswith('mps'			):
   UpperCamelCase_					:						int =					torch.manual_seed(snake_case			)
  else:
   UpperCamelCase_					:						Optional[int] =					torch.Generator(device=snake_case			).manual_seed(snake_case			)
  UpperCamelCase_					:						List[str] =					{
      '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 						SCREAMING_SNAKE_CASE__   ( self	:							List[str]   ,      snake_case	:							str   ,      snake_case	:							Dict=0			)   ->						int:

  """simple docstring"""





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

  """simple docstring"""





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

  UpperCamelCase_					:						int =					self.get_dummy_components()
  UpperCamelCase_					:						str =					self.pipeline_class(**snake_case			)
  pipe.to(snake_case			)
  pipe.set_progress_bar_config(disable=snake_case			)

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

  UpperCamelCase_					:						Any =					self.get_dummy_inputs(snake_case			)
  UpperCamelCase_					:						Any =					pipe(**snake_case			)[0]

  with tempfile.TemporaryDirectory() as tmpdir:
   pipe.save_pretrained(snake_case			)
   UpperCamelCase_					:						int =					self.pipeline_class.from_pretrained(snake_case			)
   pipe_loaded.to(snake_case			)
   pipe_loaded.set_progress_bar_config(disable=snake_case			)

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

  UpperCamelCase_					:						Dict =					self.get_dummy_inputs(snake_case			)
  UpperCamelCase_					:						int =					pipe_loaded(**snake_case			)[0]

  UpperCamelCase_					:						int =					np.abs(output - output_loaded			).max()
  self.assertLess(snake_case   ,      1e-4			)

 def 						SCREAMING_SNAKE_CASE__   ( self	:							str			)   ->						Any:

  """simple docstring"""





  UpperCamelCase_					:						Optional[int] =					'cpu'

  UpperCamelCase_					:						Any =					self.get_dummy_components()
  UpperCamelCase_					:						int =					self.pipeline_class(**snake_case			)
  pipe.to(snake_case			)
  pipe.set_progress_bar_config(disable=snake_case			)

  UpperCamelCase_					:						Union[str, Any] =					self.get_dummy_mask_inputs(snake_case			)
  UpperCamelCase_					:						Optional[int] =					pipe.generate_mask(**snake_case			)
  UpperCamelCase_					:						List[str] =					mask[0, -3:, -3:]

  self.assertEqual(mask.shape   ,      (1, 1_6, 1_6)			)
  UpperCamelCase_					:						Tuple =					np.array([0] * 9			)
  UpperCamelCase_					:						int =					np.abs(mask_slice.flatten() - expected_slice			).max()
  self.assertLessEqual(snake_case   ,      1e-3			)
  self.assertEqual(mask[0, -3, -4]   ,      0			)

 def 						SCREAMING_SNAKE_CASE__   ( self	:							Optional[Any]			)   ->						List[Any]:

  """simple docstring"""





  UpperCamelCase_					:						Dict =					'cpu'

  UpperCamelCase_					:						List[str] =					self.get_dummy_components()
  UpperCamelCase_					:						List[Any] =					self.pipeline_class(**snake_case			)
  pipe.to(snake_case			)
  pipe.set_progress_bar_config(disable=snake_case			)

  UpperCamelCase_					:						Optional[int] =					self.get_dummy_inversion_inputs(snake_case			)
  UpperCamelCase_					:						Optional[Any] =					pipe.invert(**snake_case			).images
  UpperCamelCase_					:						Tuple =					image[0, -1, -3:, -3:]

  self.assertEqual(image.shape   ,      (2, 3_2, 3_2, 3)			)
  UpperCamelCase_					:						str =					np.array(
      [0.5150, 0.5134, 0.5043, 0.5376, 0.4694, 0.51050, 0.5015, 0.4407, 0.4799]   ,      )
  UpperCamelCase_					:						Any =					np.abs(image_slice.flatten() - expected_slice			).max()
  self.assertLessEqual(snake_case   ,      1e-3			)

 def 						SCREAMING_SNAKE_CASE__   ( self	:							Any			)   ->						Any:

  """simple docstring"""





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


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

  """simple docstring"""





  UpperCamelCase_					:						Optional[Any] =					'cpu'

  UpperCamelCase_					:						Tuple =					self.get_dummy_components()

  UpperCamelCase_					:						int =					{'beta_start': 0.00085, 'beta_end': 0.012, 'beta_schedule': 'scaled_linear'}
  UpperCamelCase_					:						str =					DPMSolverMultistepScheduler(**snake_case			)
  UpperCamelCase_					:						Union[str, Any] =					DPMSolverMultistepInverseScheduler(**snake_case			)

  UpperCamelCase_					:						Optional[int] =					self.pipeline_class(**snake_case			)
  pipe.to(snake_case			)
  pipe.set_progress_bar_config(disable=snake_case			)

  UpperCamelCase_					:						Any =					self.get_dummy_inversion_inputs(snake_case			)
  UpperCamelCase_					:						Union[str, Any] =					pipe.invert(**snake_case			).images
  UpperCamelCase_					:						Any =					image[0, -1, -3:, -3:]

  self.assertEqual(image.shape   ,      (2, 3_2, 3_2, 3)			)
  UpperCamelCase_					:						int =					np.array(
      [0.5150, 0.5134, 0.5043, 0.5376, 0.4694, 0.51050, 0.5015, 0.4407, 0.4799]   ,      )
  UpperCamelCase_					:						Union[str, Any] =					np.abs(image_slice.flatten() - expected_slice			).max()
  self.assertLessEqual(snake_case   ,      1e-3			)





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

 def 						SCREAMING_SNAKE_CASE__   ( self	:							int			)   ->						Any:

  """simple docstring"""





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

 @classmethod
 def 						SCREAMING_SNAKE_CASE__   ( cls	:							Any			)   ->						Any:

  """simple docstring"""





  UpperCamelCase_					:						Optional[int] =					load_image(
      'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/diffedit/fruit.png'			)

  UpperCamelCase_					:						Any =					raw_image.convert('RGB'			).resize((7_6_8, 7_6_8)			)

  UpperCamelCase_					:						Optional[Any] =					raw_image

 def 						SCREAMING_SNAKE_CASE__   ( self	:							Optional[Any]			)   ->						Optional[int]:

  """simple docstring"""





  UpperCamelCase_					:						Dict =					torch.manual_seed(0			)

  UpperCamelCase_					:						List[str] =					StableDiffusionDiffEditPipeline.from_pretrained(
      'stabilityai/stable-diffusion-2-1'   ,      safety_checker=snake_case   ,      torch_dtype=torch.floataa			)
  UpperCamelCase_					:						int =					DDIMScheduler.from_config(pipe.scheduler.config			)
  UpperCamelCase_					:						Tuple =					DDIMInverseScheduler.from_config(pipe.scheduler.config			)
  pipe.enable_model_cpu_offload()
  pipe.set_progress_bar_config(disable=snake_case			)

  UpperCamelCase_					:						Optional[int] =					'a bowl of fruit'
  UpperCamelCase_					:						int =					'a bowl of pears'

  UpperCamelCase_					:						Any =					pipe.generate_mask(
      image=self.raw_image   ,      source_prompt=snake_case   ,      target_prompt=snake_case   ,      generator=snake_case   ,      )

  UpperCamelCase_					:						int =					pipe.invert(
      prompt=snake_case   ,      image=self.raw_image   ,      inpaint_strength=0.7   ,      generator=snake_case			).latents

  UpperCamelCase_					:						Optional[int] =					pipe(
      prompt=snake_case   ,      mask_image=snake_case   ,      image_latents=snake_case   ,      generator=snake_case   ,      negative_prompt=snake_case   ,      inpaint_strength=0.7   ,      output_type='numpy'   ,      ).images[0]

  UpperCamelCase_					:						Any =					(
      np.array(
          load_image(
              'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
              '/diffedit/pears.png'			).resize((7_6_8, 7_6_8)			)			)
      / 2_5_5
  )
  assert np.abs((expected_image - image).max()			) < 5e-1


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

  """simple docstring"""





  UpperCamelCase_					:						Union[str, Any] =					torch.manual_seed(0			)

  UpperCamelCase_					:						List[Any] =					StableDiffusionDiffEditPipeline.from_pretrained(
      'stabilityai/stable-diffusion-2-1'   ,      safety_checker=snake_case   ,      torch_dtype=torch.floataa			)
  UpperCamelCase_					:						Optional[int] =					DPMSolverMultistepScheduler.from_config(pipe.scheduler.config			)
  UpperCamelCase_					:						Union[str, Any] =					DPMSolverMultistepInverseScheduler.from_config(pipe.scheduler.config			)
  pipe.enable_model_cpu_offload()
  pipe.set_progress_bar_config(disable=snake_case			)

  UpperCamelCase_					:						str =					'a bowl of fruit'
  UpperCamelCase_					:						int =					'a bowl of pears'

  UpperCamelCase_					:						Union[str, Any] =					pipe.generate_mask(
      image=self.raw_image   ,      source_prompt=snake_case   ,      target_prompt=snake_case   ,      generator=snake_case   ,      )

  UpperCamelCase_					:						Any =					pipe.invert(
      prompt=snake_case   ,      image=self.raw_image   ,      inpaint_strength=0.7   ,      generator=snake_case   ,      num_inference_steps=2_5   ,      ).latents

  UpperCamelCase_					:						int =					pipe(
      prompt=snake_case   ,      mask_image=snake_case   ,      image_latents=snake_case   ,      generator=snake_case   ,      negative_prompt=snake_case   ,      inpaint_strength=0.7   ,      num_inference_steps=2_5   ,      output_type='numpy'   ,      ).images[0]

  UpperCamelCase_					:						str =					(
      np.array(
          load_image(
              'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
              '/diffedit/pears.png'			).resize((7_6_8, 7_6_8)			)			)
      / 2_5_5
  )
  assert np.abs((expected_image - image).max()			) < 5e-1



