


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    UpperCamelCase    (						UpperCamelCase__  ,					UpperCamelCase__  ,					UpperCamelCase__  ,					unittest.TestCase   ):


   '''simple docstring'''







   lowercase    :		Tuple	  =StableDiffusionModelEditingPipeline
   lowercase    :		Dict	  =TEXT_TO_IMAGE_PARAMS
   lowercase    :		Dict	  =TEXT_TO_IMAGE_BATCH_PARAMS
   lowercase    :		str	  =TEXT_TO_IMAGE_IMAGE_PARAMS
   lowercase    :		Optional[Any]	  =TEXT_TO_IMAGE_IMAGE_PARAMS
   def 						UpperCamelCase     (   self						):
        torch.manual_seed(0						)
        lowercase_				:Any       			=				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      ,							)
        lowercase_				:Optional[int]       			=				DDIMScheduler()
        torch.manual_seed(0						)
        lowercase_				:Optional[int]       			=				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						)
        lowercase_				:str       			=				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      ,							)
        lowercase_				:str       			=				CLIPTextModel(SCREAMING_SNAKE_CASE__						)
        lowercase_				:Dict       			=				CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''						)

        lowercase_				:str       			=				{
            '''unet''': unet,
            '''scheduler''': scheduler,
            '''vae''': vae,
            '''text_encoder''': text_encoder,
            '''tokenizer''': tokenizer,
            '''safety_checker''': None,
            '''feature_extractor''': None,
        }
        return components
   def 						UpperCamelCase     (   self      ,							UpperCamelCase_      ,							UpperCamelCase_=0						):
        lowercase_				:Optional[Any]       			=				torch.manual_seed(SCREAMING_SNAKE_CASE__						)
        lowercase_				:int       			=				{
            '''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 						UpperCamelCase     (   self						):
        lowercase_				:Any       			=				'''cpu'''  # ensure determinism for the device-dependent torch.Generator
        lowercase_				:Dict       			=				self.get_dummy_components()
        lowercase_				:List[str]       			=				StableDiffusionModelEditingPipeline(**SCREAMING_SNAKE_CASE__						)
        lowercase_				:Optional[int]       			=				sd_pipe.to(SCREAMING_SNAKE_CASE__						)
        sd_pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__						)

        lowercase_				:List[Any]       			=				self.get_dummy_inputs(SCREAMING_SNAKE_CASE__						)
        lowercase_				:int       			=				sd_pipe(**SCREAMING_SNAKE_CASE__						).images
        lowercase_				:int       			=				image[0, -3:, -3:, -1]
        assert image.shape == (1, 64, 64, 3)

        lowercase_				:List[Any]       			=				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 						UpperCamelCase     (   self						):
        lowercase_				:int       			=				'''cpu'''  # ensure determinism for the device-dependent torch.Generator
        lowercase_				:List[str]       			=				self.get_dummy_components()
        lowercase_				:List[str]       			=				StableDiffusionModelEditingPipeline(**SCREAMING_SNAKE_CASE__						)
        lowercase_				:List[Any]       			=				sd_pipe.to(SCREAMING_SNAKE_CASE__						)
        sd_pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__						)

        lowercase_				:Any       			=				self.get_dummy_inputs(SCREAMING_SNAKE_CASE__						)
        lowercase_				:Union[str, Any]       			=				'''french fries'''
        lowercase_				:Any       			=				sd_pipe(**SCREAMING_SNAKE_CASE__      ,							negative_prompt=SCREAMING_SNAKE_CASE__						)
        lowercase_				:Optional[Any]       			=				output.images
        lowercase_				:Optional[int]       			=				image[0, -3:, -3:, -1]

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

        lowercase_				:int       			=				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 						UpperCamelCase     (   self						):
        lowercase_				:Any       			=				'''cpu'''  # ensure determinism for the device-dependent torch.Generator
        lowercase_				:int       			=				self.get_dummy_components()
        lowercase_				:Tuple       			=				EulerAncestralDiscreteScheduler(
            beta_start=0.0_0085      ,							beta_end=0.012      ,							beta_schedule='''scaled_linear'''						)
        lowercase_				:Optional[Any]       			=				StableDiffusionModelEditingPipeline(**SCREAMING_SNAKE_CASE__						)
        lowercase_				:str       			=				sd_pipe.to(SCREAMING_SNAKE_CASE__						)
        sd_pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__						)

        lowercase_				:Optional[Any]       			=				self.get_dummy_inputs(SCREAMING_SNAKE_CASE__						)
        lowercase_				:str       			=				sd_pipe(**SCREAMING_SNAKE_CASE__						).images
        lowercase_				:Tuple       			=				image[0, -3:, -3:, -1]

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

        lowercase_				:Dict       			=				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 						UpperCamelCase     (   self						):
        lowercase_				:Any       			=				'''cpu'''  # ensure determinism for the device-dependent torch.Generator
        lowercase_				:Optional[Any]       			=				self.get_dummy_components()
        lowercase_				:Tuple       			=				PNDMScheduler()
        lowercase_				:int       			=				StableDiffusionModelEditingPipeline(**SCREAMING_SNAKE_CASE__						)
        lowercase_				:List[Any]       			=				sd_pipe.to(SCREAMING_SNAKE_CASE__						)
        sd_pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__						)

        lowercase_				:List[Any]       			=				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__						):
             lowercase_				:Dict       			=				sd_pipe(**SCREAMING_SNAKE_CASE__						).images
   def 						UpperCamelCase     (   self						):
        super().test_inference_batch_single_identical(expected_max_diff=5E-3						)




   def 						UpperCamelCase     (   self						):
        super().test_attention_slicing_forward_pass(expected_max_diff=5E-3						)






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


   '''simple docstring'''







   def 						UpperCamelCase     (   self						):
        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()
   def 						UpperCamelCase     (   self      ,							UpperCamelCase_=0						):
        lowercase_				:List[str]       			=				torch.manual_seed(SCREAMING_SNAKE_CASE__						)
        lowercase_				:Optional[int]       			=				{
            '''prompt''': '''A field of roses''',
            '''generator''': generator,
            '''num_inference_steps''': 3,
            '''guidance_scale''': 7.5,
            '''output_type''': '''numpy''',
        }
        return inputs
   def 						UpperCamelCase     (   self						):
        lowercase_				:int       			=				'''CompVis/stable-diffusion-v1-4'''
        lowercase_				:Union[str, Any]       			=				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()

        lowercase_				:Dict       			=				self.get_inputs()
        lowercase_				:List[Any]       			=				pipe(**SCREAMING_SNAKE_CASE__						).images
        lowercase_				:Dict       			=				image[0, -3:, -3:, -1].flatten()

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

        lowercase_				:int       			=				np.array(
            [0.674_9496, 0.638_6453, 0.5144_3267, 0.6609_4905, 0.6192_1215, 0.549_1332, 0.574_4417, 0.5807_5106, 0.517_4658]						)

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

        lowercase_				:str       			=				pipe(**SCREAMING_SNAKE_CASE__						).images
        lowercase_				:Optional[Any]       			=				image[0, -3:, -3:, -1].flatten()

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

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




   def 						UpperCamelCase     (   self						):
        torch.cuda.empty_cache()
        torch.cuda.reset_max_memory_allocated()
        torch.cuda.reset_peak_memory_stats()

        lowercase_				:List[str]       			=				'''CompVis/stable-diffusion-v1-4'''
        lowercase_				:str       			=				DDIMScheduler.from_pretrained(SCREAMING_SNAKE_CASE__      ,							subfolder='''scheduler'''						)
        lowercase_				:Tuple       			=				StableDiffusionModelEditingPipeline.from_pretrained(
            SCREAMING_SNAKE_CASE__      ,							scheduler=SCREAMING_SNAKE_CASE__      ,							safety_checker=SCREAMING_SNAKE_CASE__						)
        lowercase_				:int       			=				pipe.to(SCREAMING_SNAKE_CASE__						)
        pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__						)
        pipe.enable_attention_slicing(1						)
        pipe.enable_sequential_cpu_offload()

        lowercase_				:Union[str, Any]       			=				self.get_inputs()
        lowercase_				:Tuple       			=				pipe(**SCREAMING_SNAKE_CASE__						)

        lowercase_				:Union[str, Any]       			=				torch.cuda.max_memory_allocated()
        # make sure that less than 4.4 GB is allocated
        assert mem_bytes < 4.4 * 10**9







# limitations under the License.


from typing import Optional, Tuple, Union

import torch

from diffusers.pipeline_utils import DiffusionPipeline, ImagePipelineOutput





class  __lowerCamelCase       (    UpperCamelCase__    ):




    """simple docstring"""

    def __init__(       self     :     Union[str, Any]   ,	SCREAMING_SNAKE_CASE__     :     Union[str, Any]   ,	SCREAMING_SNAKE_CASE__     :     Any   )      ->       Union[str, Any]:
         super().__init__()
         self.register_modules(unet=SCREAMING_SNAKE_CASE__   ,	scheduler=SCREAMING_SNAKE_CASE__   )


    @torch.no_grad()
    def __call__(       self     :     Dict   ,	SCREAMING_SNAKE_CASE__     :     int = 1   ,	SCREAMING_SNAKE_CASE__     :     Optional[torch.Generator] = None   ,	SCREAMING_SNAKE_CASE__     :     int = 50   ,	SCREAMING_SNAKE_CASE__     :     Optional[str] = "pil"   ,	SCREAMING_SNAKE_CASE__     :     bool = True   ,	**SCREAMING_SNAKE_CASE__     :     Union[str, Any]   ,	)      ->       Union[ImagePipelineOutput, Tuple]:
         lowerCAmelCase__		= torch.randn(
             (batch_size, self.unet.config.in_channels, self.unet.config.sample_size, self.unet.config.sample_size)   ,	generator=SCREAMING_SNAKE_CASE__   ,	)
         lowerCAmelCase__		= image.to(self.device   )

         # set step values
         self.scheduler.set_timesteps(SCREAMING_SNAKE_CASE__   )

         for t in self.progress_bar(self.scheduler.timesteps   ):
              # 1. predict noise model_output
              lowerCAmelCase__		= self.unet(SCREAMING_SNAKE_CASE__   ,	SCREAMING_SNAKE_CASE__   ).sample

              # 2. predict previous mean of image x_t-1 and add variance depending on eta
              # eta corresponds to η in paper and should be between [0, 1]
              # do x_t -> x_t-1
              lowerCAmelCase__		= self.scheduler.step(SCREAMING_SNAKE_CASE__   ,	SCREAMING_SNAKE_CASE__   ,	SCREAMING_SNAKE_CASE__   ).prev_sample

         lowerCAmelCase__		= (image / 2 + 0.5).clamp(0   ,	1   )
         lowerCAmelCase__		= image.cpu().permute(0   ,	2   ,	3   ,	1   ).numpy()
         if output_type == "pil":
              lowerCAmelCase__		= self.numpy_to_pil(SCREAMING_SNAKE_CASE__   )

         if not return_dict:
              return (image,), "This is a local test"

         return ImagePipelineOutput(images=SCREAMING_SNAKE_CASE__   ), "This is a local test"
