



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





 """simple docstring"""



 lowerCamelCase_       =     StableDiffusionModelEditingPipeline
 lowerCamelCase_       =     TEXT_TO_IMAGE_PARAMS
 lowerCamelCase_       =     TEXT_TO_IMAGE_BATCH_PARAMS
 lowerCamelCase_       =     TEXT_TO_IMAGE_IMAGE_PARAMS
 lowerCamelCase_       =     TEXT_TO_IMAGE_IMAGE_PARAMS
 def lowercase_						(  self   )		->							Any:




   """simple docstring"""







   torch.manual_seed(0   )
   _A          =      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		,      )
   _A          =      DDIMScheduler()
   torch.manual_seed(0   )
   _A          =      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		,      )
   torch.manual_seed(0   )
   _A          =      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		,      )
   _A          =      CLIPTextModel(a   )
   _A          =      CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''   )

   _A          =      {
       '''unet''': unet,
       '''scheduler''': scheduler,
       '''vae''': vae,
       '''text_encoder''': text_encoder,
       '''tokenizer''': tokenizer,
       '''safety_checker''': None,
       '''feature_extractor''': None,
   }
   return components
 def lowercase_						(  self		,      a		,      a=0   )		->							List[str]:




   """simple docstring"""







   _A          =      torch.manual_seed(a   )
   _A          =      {
       '''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 lowercase_						(  self   )		->							Optional[int]:




   """simple docstring"""







   _A          =      '''cpu'''  # ensure determinism for the device-dependent torch.Generator
   _A          =      self.get_dummy_components()
   _A          =      StableDiffusionModelEditingPipeline(**a   )
   _A          =      sd_pipe.to(a   )
   sd_pipe.set_progress_bar_config(disable=a   )

   _A          =      self.get_dummy_inputs(a   )
   _A          =      sd_pipe(**a   ).images
   _A          =      image[0, -3:, -3:, -1]
   assert image.shape == (1, 6_4, 6_4, 3)

   _A          =      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 lowercase_						(  self   )		->							Any:




   """simple docstring"""







   _A          =      '''cpu'''  # ensure determinism for the device-dependent torch.Generator
   _A          =      self.get_dummy_components()
   _A          =      StableDiffusionModelEditingPipeline(**a   )
   _A          =      sd_pipe.to(a   )
   sd_pipe.set_progress_bar_config(disable=a   )

   _A          =      self.get_dummy_inputs(a   )
   _A          =      '''french fries'''
   _A          =      sd_pipe(**a		,      negative_prompt=a   )
   _A          =      output.images
   _A          =      image[0, -3:, -3:, -1]

   assert image.shape == (1, 6_4, 6_4, 3)

   _A          =      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 lowercase_						(  self   )		->							Optional[int]:




   """simple docstring"""







   _A          =      '''cpu'''  # ensure determinism for the device-dependent torch.Generator
   _A          =      self.get_dummy_components()
   _A          =      EulerAncestralDiscreteScheduler(
       beta_start=0.00_085		,      beta_end=0.012		,      beta_schedule='''scaled_linear'''   )
   _A          =      StableDiffusionModelEditingPipeline(**a   )
   _A          =      sd_pipe.to(a   )
   sd_pipe.set_progress_bar_config(disable=a   )

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

   assert image.shape == (1, 6_4, 6_4, 3)

   _A          =      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 lowercase_						(  self   )		->							Optional[Any]:




   """simple docstring"""







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

   _A          =      self.get_dummy_inputs(a   )
   # the pipeline does not expect pndm so test if it raises error.
   with self.assertRaises(a   ):
     _A          =      sd_pipe(**a   ).images
 def lowercase_						(  self   )		->							str:




   """simple docstring"""







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



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




   """simple docstring"""







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







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





 """simple docstring"""



 def lowercase_						(  self   )		->							Any:




   """simple docstring"""







   super().tearDown()
   gc.collect()
   torch.cuda.empty_cache()
 def lowercase_						(  self		,      a=0   )		->							Optional[Any]:




   """simple docstring"""







   _A          =      torch.manual_seed(a   )
   _A          =      {
       '''prompt''': '''A field of roses''',
       '''generator''': generator,
       '''num_inference_steps''': 3,
       '''guidance_scale''': 7.5,
       '''output_type''': '''numpy''',
   }
   return inputs
 def lowercase_						(  self   )		->							List[str]:




   """simple docstring"""







   _A          =      '''CompVis/stable-diffusion-v1-4'''
   _A          =      StableDiffusionModelEditingPipeline.from_pretrained(a		,      safety_checker=a   )
   pipe.to(a   )
   pipe.set_progress_bar_config(disable=a   )
   pipe.enable_attention_slicing()

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

   assert image.shape == (1, 5_1_2, 5_1_2, 3)

   _A          =      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'''   )

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

   assert image.shape == (1, 5_1_2, 5_1_2, 3)

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



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




   """simple docstring"""







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

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

   _A          =      self.get_inputs()
   _A          =      pipe(**a   )

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




import importlib.metadata
import subprocess
import sys




def       UpperCAmelCase__	(							__snake_case = False							)				->	List[Any]:
  _A          =      False
  if "IPython" in sys.modules:
    _A          =      '''google.colab''' in str(sys.modules['''IPython'''].get_ipython()							)

  if in_colab:
    if upgrade:
      _A          =      ['''pip''', '''install''', '''-U''', '''torch''']
      subprocess.run(__snake_case  ,			check=__snake_case							)
    # get the current version of torch
    _A          =      importlib.metadata.version('''torch'''							)
    _A          =      torch_version[: torch_version.rindex('''.'''							)]
    _A          =      F'''https://storage.googleapis.com/tpu-pytorch/wheels/colab/torch_xla-{torch_version_trunc}-cp37-cp37m-linux_x86_64.whl'''
    _A          =      ['''pip''', '''install''', xla_wheel]
    subprocess.run(__snake_case  ,			check=__snake_case							)
  else:
    raise RuntimeError('''`install_xla` utility works only on google colab.'''							)