
'''simple docstring'''







from logging import getLogger
from typing import Any, Callable, List, Optional, Union

import numpy as np
import PIL
import torch

from ...schedulers import DDPMScheduler
from ..onnx_utils import ORT_TO_NP_TYPE, OnnxRuntimeModel
from ..pipeline_utils import ImagePipelineOutput
from . import StableDiffusionUpscalePipeline


lowerCAmelCase_      =   getLogger(__name__)


lowerCAmelCase_      =   4
lowerCAmelCase_      =   7

lowerCAmelCase_      =   {
    "float16": torch.floataa,
    "float32": torch.floataa,
}







def   lowerCAmelCase(  a__    :  Dict      ):







       '''simple docstring'''






       if isinstance(a__      , torch.Tensor      ):
              return image
       elif isinstance(a__      , PIL.Image.Image      ):
              lowerCamelCase__				   =      [image]

       if isinstance(image[0]      , PIL.Image.Image      ):
              lowerCamelCase__   ,		lowerCamelCase__				   =      image[0].size
              lowerCamelCase__   ,		lowerCamelCase__				   =      (x - x % 64 for x in (w, h))  # resize to integer multiple of 32

              lowerCamelCase__				   =      [np.array(i.resize((w, h)      )      )[None, :] for i in image]
              lowerCamelCase__				   =      np.concatenate(a__      , axis=0      )
              lowerCamelCase__				   =      np.array(a__      ).astype(np.floataa      ) / 2_5_5.0
              lowerCamelCase__				   =      image.transpose(0      , 3      , 1      , 2      )
              lowerCamelCase__				   =      2.0 * image - 1.0
              lowerCamelCase__				   =      torch.from_numpy(a__      )
       elif isinstance(image[0]      , torch.Tensor      ):
              lowerCamelCase__				   =      torch.cat(a__      , dim=0      )

       return image







class 		snake_case_       (						A__						):



   """simple docstring"""





   def __init__( self			,   UpperCamelCase			,   UpperCamelCase			,   UpperCamelCase			,   UpperCamelCase			,   UpperCamelCase			,   UpperCamelCase			,   UpperCamelCase = 3_50			,   ):
          super().__init__(
              vae=UpperCamelCase			,   text_encoder=UpperCamelCase			,   tokenizer=UpperCamelCase			,   unet=UpperCamelCase			,   low_res_scheduler=UpperCamelCase			,   scheduler=UpperCamelCase			,   safety_checker=UpperCamelCase			,   feature_extractor=UpperCamelCase			,   watermarker=UpperCamelCase			,   max_noise_level=UpperCamelCase			,   )







   def __call__( self			,   UpperCamelCase			,   UpperCamelCase			,   UpperCamelCase = 75			,   UpperCamelCase = 9.0			,   UpperCamelCase = 20			,   UpperCamelCase = None			,   UpperCamelCase = 1			,   UpperCamelCase = 0.0			,   UpperCamelCase = None			,   UpperCamelCase = None			,   UpperCamelCase = None			,   UpperCamelCase = None			,   UpperCamelCase = "pil"			,   UpperCamelCase = True			,   UpperCamelCase = None			,   UpperCamelCase = 1			,   ):
          self.check_inputs(UpperCamelCase			,   UpperCamelCase			,   UpperCamelCase			,   UpperCamelCase)

          # 2. Define call parameters
          if prompt is not None and isinstance(UpperCamelCase			,   UpperCamelCase):
                 lowerCamelCase__				   =      1
          elif prompt is not None and isinstance(UpperCamelCase			,   UpperCamelCase):
                 lowerCamelCase__				   =      len(UpperCamelCase)
          else:
                 lowerCamelCase__				   =      prompt_embeds.shape[0]

          lowerCamelCase__				   =      self._execution_device
          # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
          # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
          # corresponds to doing no classifier free guidance.
          lowerCamelCase__				   =      guidance_scale > 1.0

          # 3. Encode input prompt
          lowerCamelCase__				   =      self._encode_prompt(
              UpperCamelCase			,   UpperCamelCase			,   UpperCamelCase			,   UpperCamelCase			,   UpperCamelCase			,   prompt_embeds=UpperCamelCase			,   negative_prompt_embeds=UpperCamelCase			,   )

          lowerCamelCase__				   =      ORT_TO_PT_TYPE[str(text_embeddings.dtype)]

          # 4. Preprocess image
          lowerCamelCase__				   =      preprocess(UpperCamelCase)
          lowerCamelCase__				   =      image.cpu()

          # 5. set timesteps
          self.scheduler.set_timesteps(UpperCamelCase			,   device=UpperCamelCase)
          lowerCamelCase__				   =      self.scheduler.timesteps

          # 5. Add noise to image
          lowerCamelCase__				   =      torch.tensor([noise_level]			,   dtype=torch.long			,   device=UpperCamelCase)
          lowerCamelCase__				   =      torch.randn(image.shape			,   generator=UpperCamelCase			,   device=UpperCamelCase			,   dtype=UpperCamelCase)
          lowerCamelCase__				   =      self.low_res_scheduler.add_noise(UpperCamelCase			,   UpperCamelCase			,   UpperCamelCase)

          lowerCamelCase__				   =      2 if do_classifier_free_guidance else 1
          lowerCamelCase__				   =      np.concatenate([image] * batch_multiplier * num_images_per_prompt)
          lowerCamelCase__				   =      np.concatenate([noise_level] * image.shape[0])

          # 6. Prepare latent variables
          lowerCamelCase__   ,		lowerCamelCase__				   =      image.shape[2:]
          lowerCamelCase__				   =      self.prepare_latents(
              batch_size * num_images_per_prompt			,   UpperCamelCase			,   UpperCamelCase			,   UpperCamelCase			,   UpperCamelCase			,   UpperCamelCase			,   UpperCamelCase			,   UpperCamelCase			,   )

          # 7. Check that sizes of image and latents match
          lowerCamelCase__				   =      image.shape[1]
          if NUM_LATENT_CHANNELS + num_channels_image != NUM_UNET_INPUT_CHANNELS:
                 raise ValueError(
                     "Incorrect configuration settings! The config of `pipeline.unet` expects"
                     f""" {NUM_UNET_INPUT_CHANNELS} but received `num_channels_latents`: {NUM_LATENT_CHANNELS} +"""
                     f""" `num_channels_image`: {num_channels_image} """
                     f""" = {NUM_LATENT_CHANNELS+num_channels_image}. Please verify the config of"""
                     " `pipeline.unet` or your `image` input.")

          # 8. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
          lowerCamelCase__				   =      self.prepare_extra_step_kwargs(UpperCamelCase			,   UpperCamelCase)

          lowerCamelCase__				   =      next(
              (input.type for input in self.unet.model.get_inputs() if input.name == "timestep")			,   "tensor(float)")
          lowerCamelCase__				   =      ORT_TO_NP_TYPE[timestep_dtype]

          # 9. Denoising loop
          lowerCamelCase__				   =      len(UpperCamelCase) - num_inference_steps * self.scheduler.order
          with self.progress_bar(total=UpperCamelCase) as progress_bar:
                 for i, t in enumerate(UpperCamelCase):
                        # expand the latents if we are doing classifier free guidance
                        lowerCamelCase__				   =      np.concatenate([latents] * 2) if do_classifier_free_guidance else latents

                        # concat latents, mask, masked_image_latents in the channel dimension
                        lowerCamelCase__				   =      self.scheduler.scale_model_input(UpperCamelCase			,   UpperCamelCase)
                        lowerCamelCase__				   =      np.concatenate([latent_model_input, image]			,   axis=1)

                        # timestep to tensor
                        lowerCamelCase__				   =      np.array([t]			,   dtype=UpperCamelCase)

                        # predict the noise residual
                        lowerCamelCase__				   =      self.unet(
                            sample=UpperCamelCase			,   timestep=UpperCamelCase			,   encoder_hidden_states=UpperCamelCase			,   class_labels=noise_level.astype(np.intaa)			,   )[0]

                        # perform guidance
                        if do_classifier_free_guidance:
                               lowerCamelCase__   ,		lowerCamelCase__				   =      np.split(UpperCamelCase			,   2)
                        lowerCamelCase__				   =      noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

                        # compute the previous noisy sample x_t -> x_t-1
                        lowerCamelCase__				   =      self.scheduler.step(
                            torch.from_numpy(UpperCamelCase)			,   UpperCamelCase			,   UpperCamelCase			,   **UpperCamelCase).prev_sample

                        # call the callback, if provided
                        if i == len(UpperCamelCase) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
                               progress_bar.update()
                               if callback is not None and i % callback_steps == 0:
                                      callback(UpperCamelCase			,   UpperCamelCase			,   UpperCamelCase)

        # 10. Post-processing
          lowerCamelCase__				   =      self.decode_latents(latents.float())

          # 11. Convert to PIL
          if output_type == "pil":
                 lowerCamelCase__				   =      self.numpy_to_pil(UpperCamelCase)

          if not return_dict:
                 return (image,)

          return ImagePipelineOutput(images=UpperCamelCase)







   def     __UpperCAmelCase      ( self			,   UpperCamelCase):
          lowerCamelCase__				   =      1 / 0.0_8_3_3_3 * latents
          lowerCamelCase__				   =      self.vae(latent_sample=UpperCamelCase)[0]
          lowerCamelCase__				   =      np.clip(image / 2 + 0.5			,   0			,   1)
          lowerCamelCase__				   =      image.transpose((0, 2, 3, 1))
          return image







   def     __UpperCAmelCase      ( self			,   UpperCamelCase			,   UpperCamelCase			,   UpperCamelCase			,   UpperCamelCase			,   UpperCamelCase			,   UpperCamelCase = None			,   UpperCamelCase = None			,   ):
          if prompt is not None and isinstance(UpperCamelCase			,   UpperCamelCase):
                 lowerCamelCase__				   =      1
          elif prompt is not None and isinstance(UpperCamelCase			,   UpperCamelCase):
                 lowerCamelCase__				   =      len(UpperCamelCase)
          else:
                 lowerCamelCase__				   =      prompt_embeds.shape[0]

          if prompt_embeds is None:
                 lowerCamelCase__				   =      self.tokenizer(
                     UpperCamelCase			,   padding="max_length"			,   max_length=self.tokenizer.model_max_length			,   truncation=UpperCamelCase			,   return_tensors="pt"			,   )
                 lowerCamelCase__				   =      text_inputs.input_ids
                 lowerCamelCase__				   =      self.tokenizer(UpperCamelCase			,   padding="longest"			,   return_tensors="pt").input_ids

                 if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(
                     UpperCamelCase			,   UpperCamelCase):
                        lowerCamelCase__				   =      self.tokenizer.batch_decode(
                            untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1])
                        logger.warning(
                            "The following part of your input was truncated because CLIP can only handle sequences up to"
                            f""" {self.tokenizer.model_max_length} tokens: {removed_text}""")

                 # no positional arguments to text_encoder
                 lowerCamelCase__				   =      self.text_encoder(
                     input_ids=text_input_ids.int().to(UpperCamelCase)			,   )
                 lowerCamelCase__				   =      prompt_embeds[0]

          lowerCamelCase__   ,		lowerCamelCase__   ,		lowerCamelCase__				   =      prompt_embeds.shape
          # duplicate text embeddings for each generation per prompt, using mps friendly method
          lowerCamelCase__				   =      prompt_embeds.repeat(1			,   UpperCamelCase)
          lowerCamelCase__				   =      prompt_embeds.reshape(bs_embed * num_images_per_prompt			,   UpperCamelCase			,   -1)

          # get unconditional embeddings for classifier free guidance
          if do_classifier_free_guidance and negative_prompt_embeds is None:
                 lowerCamelCase__				   =      42
                 if negative_prompt is None:
                        lowerCamelCase__				   =      [""] * batch_size
                 elif type(UpperCamelCase) is not type(UpperCamelCase):
                        raise TypeError(
                            f"""`negative_prompt` should be the same type to `prompt`, but got {type(UpperCamelCase)} !="""
                            f""" {type(UpperCamelCase)}.""")
                 elif isinstance(UpperCamelCase			,   UpperCamelCase):
                        lowerCamelCase__				   =      [negative_prompt]
                 elif batch_size != len(UpperCamelCase):
                        raise ValueError(
                            f"""`negative_prompt`: {negative_prompt} has batch size {len(UpperCamelCase)}, but `prompt`:"""
                            f""" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"""
                            " the batch size of `prompt`.")
                 else:
                        lowerCamelCase__				   =      negative_prompt

                 lowerCamelCase__				   =      text_input_ids.shape[-1]
                 lowerCamelCase__				   =      self.tokenizer(
                     UpperCamelCase			,   padding="max_length"			,   max_length=UpperCamelCase			,   truncation=UpperCamelCase			,   return_tensors="pt"			,   )

                 # if hasattr(uncond_input, "attention_mask"):
                 #     attention_mask = uncond_input.attention_mask.to(device)
                 # else:
                 #     attention_mask = None

                 lowerCamelCase__				   =      self.text_encoder(
                     input_ids=uncond_input.input_ids.int().to(UpperCamelCase)			,   )
                 lowerCamelCase__				   =      uncond_embeddings[0]

          if do_classifier_free_guidance:
                 lowerCamelCase__				   =      uncond_embeddings.shape[1]
                 # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
                 lowerCamelCase__				   =      uncond_embeddings.repeat(1			,   UpperCamelCase)
                 lowerCamelCase__				   =      uncond_embeddings.reshape(batch_size * num_images_per_prompt			,   UpperCamelCase			,   -1)

                 # For classifier free guidance, we need to do two forward passes.
                 # Here we concatenate the unconditional and text embeddings into a single batch
                 # to avoid doing two forward passes
                 lowerCamelCase__				   =      np.concatenate([uncond_embeddings, prompt_embeds])

          return prompt_embeds



'''simple docstring'''







import gc
import random
import tempfile
import unittest

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

from diffusers import AutoencoderKL, DDIMScheduler, LMSDiscreteScheduler, PNDMScheduler, UNetaDConditionModel
from diffusers.pipelines.semantic_stable_diffusion import SemanticStableDiffusionPipeline as StableDiffusionPipeline
from diffusers.utils import floats_tensor, nightly, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()







class 		snake_case_       (						unittest.TestCase						):



   """simple docstring"""





   def     __UpperCAmelCase      ( self):
          # clean up the VRAM after each test
          super().tearDown()
          gc.collect()
          torch.cuda.empty_cache()







   @property
   def     __UpperCAmelCase      ( self):
          lowerCamelCase__				   =      1
          lowerCamelCase__				   =      3
          lowerCamelCase__				   =      (32, 32)

          lowerCamelCase__				   =      floats_tensor((batch_size, num_channels) + sizes			,   rng=random.Random(0)).to(UpperCamelCase)
          return image







   @property
   def     __UpperCAmelCase      ( self):
          torch.manual_seed(0)
          lowerCamelCase__				   =      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			,   )
          return model







   @property
   def     __UpperCAmelCase      ( self):
          torch.manual_seed(0)
          lowerCamelCase__				   =      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			,   )
          return model







   @property
   def     __UpperCAmelCase      ( self):
          torch.manual_seed(0)
          lowerCamelCase__				   =      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=10_00			,   )
          return CLIPTextModel(UpperCamelCase)







   @property
   def     __UpperCAmelCase      ( self):
          def extract(*UpperCamelCase			,   **UpperCamelCase):







                 class 		snake_case_       :



                    """simple docstring"""





                    def __init__( self):
                           lowerCamelCase__				   =      torch.ones([0])







                    def     __UpperCAmelCase      ( self			,   UpperCamelCase):
                           self.pixel_values.to(UpperCamelCase)
                           return self






                 return Out()

          return extract







   def     __UpperCAmelCase      ( self):
          lowerCamelCase__				   =      "cpu"  # ensure determinism for the device-dependent torch.Generator
          lowerCamelCase__				   =      self.dummy_cond_unet
          lowerCamelCase__				   =      DDIMScheduler(
              beta_start=0.0_0_0_8_5			,   beta_end=0.0_1_2			,   beta_schedule="scaled_linear"			,   clip_sample=UpperCamelCase			,   set_alpha_to_one=UpperCamelCase			,   )

          lowerCamelCase__				   =      self.dummy_vae
          lowerCamelCase__				   =      self.dummy_text_encoder
          lowerCamelCase__				   =      CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")

          # make sure here that pndm scheduler skips prk
          lowerCamelCase__				   =      StableDiffusionPipeline(
              unet=UpperCamelCase			,   scheduler=UpperCamelCase			,   vae=UpperCamelCase			,   text_encoder=UpperCamelCase			,   tokenizer=UpperCamelCase			,   safety_checker=UpperCamelCase			,   feature_extractor=self.dummy_extractor			,   )
          lowerCamelCase__				   =      sd_pipe.to(UpperCamelCase)
          sd_pipe.set_progress_bar_config(disable=UpperCamelCase)

          lowerCamelCase__				   =      "A painting of a squirrel eating a burger"

          lowerCamelCase__				   =      torch.Generator(device=UpperCamelCase).manual_seed(0)
          lowerCamelCase__				   =      sd_pipe([prompt]			,   generator=UpperCamelCase			,   guidance_scale=6.0			,   num_inference_steps=2			,   output_type="np")
          lowerCamelCase__				   =      output.images

          lowerCamelCase__				   =      torch.Generator(device=UpperCamelCase).manual_seed(0)
          lowerCamelCase__				   =      sd_pipe(
              [prompt]			,   generator=UpperCamelCase			,   guidance_scale=6.0			,   num_inference_steps=2			,   output_type="np"			,   return_dict=UpperCamelCase			,   )[0]

          lowerCamelCase__				   =      image[0, -3:, -3:, -1]
          lowerCamelCase__				   =      image_from_tuple[0, -3:, -3:, -1]

          assert image.shape == (1, 64, 64, 3)
          lowerCamelCase__				   =      np.array([0.5_7_5_3, 0.6_1_1_4, 0.5_0_0_1, 0.5_0_3_4, 0.5_4_7_0, 0.4_7_2_9, 0.4_9_7_1, 0.4_8_6_7, 0.4_8_6_7])

          assert np.abs(image_slice.flatten() - expected_slice).max() < 1E-2
          assert np.abs(image_from_tuple_slice.flatten() - expected_slice).max() < 1E-2







   def     __UpperCAmelCase      ( self):
          lowerCamelCase__				   =      "cpu"  # ensure determinism for the device-dependent torch.Generator
          lowerCamelCase__				   =      self.dummy_cond_unet
          lowerCamelCase__				   =      PNDMScheduler(skip_prk_steps=UpperCamelCase)
          lowerCamelCase__				   =      self.dummy_vae
          lowerCamelCase__				   =      self.dummy_text_encoder
          lowerCamelCase__				   =      CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")

          # make sure here that pndm scheduler skips prk
          lowerCamelCase__				   =      StableDiffusionPipeline(
              unet=UpperCamelCase			,   scheduler=UpperCamelCase			,   vae=UpperCamelCase			,   text_encoder=UpperCamelCase			,   tokenizer=UpperCamelCase			,   safety_checker=UpperCamelCase			,   feature_extractor=self.dummy_extractor			,   )
          lowerCamelCase__				   =      sd_pipe.to(UpperCamelCase)
          sd_pipe.set_progress_bar_config(disable=UpperCamelCase)

          lowerCamelCase__				   =      "A painting of a squirrel eating a burger"
          lowerCamelCase__				   =      torch.Generator(device=UpperCamelCase).manual_seed(0)
          lowerCamelCase__				   =      sd_pipe([prompt]			,   generator=UpperCamelCase			,   guidance_scale=6.0			,   num_inference_steps=2			,   output_type="np")

          lowerCamelCase__				   =      output.images

          lowerCamelCase__				   =      torch.Generator(device=UpperCamelCase).manual_seed(0)
          lowerCamelCase__				   =      sd_pipe(
              [prompt]			,   generator=UpperCamelCase			,   guidance_scale=6.0			,   num_inference_steps=2			,   output_type="np"			,   return_dict=UpperCamelCase			,   )[0]

          lowerCamelCase__				   =      image[0, -3:, -3:, -1]
          lowerCamelCase__				   =      image_from_tuple[0, -3:, -3:, -1]

          assert image.shape == (1, 64, 64, 3)
          lowerCamelCase__				   =      np.array([0.5_1_2_2, 0.5_7_1_2, 0.4_8_2_5, 0.5_0_5_3, 0.5_6_4_6, 0.4_7_6_9, 0.5_1_7_9, 0.4_8_9_4, 0.4_9_9_4])

          assert np.abs(image_slice.flatten() - expected_slice).max() < 1E-2
          assert np.abs(image_from_tuple_slice.flatten() - expected_slice).max() < 1E-2







   def     __UpperCAmelCase      ( self):
          lowerCamelCase__				   =      StableDiffusionPipeline.from_pretrained(
              "hf-internal-testing/tiny-stable-diffusion-lms-pipe"			,   safety_checker=UpperCamelCase)
          assert isinstance(UpperCamelCase			,   UpperCamelCase)
          assert isinstance(pipe.scheduler			,   UpperCamelCase)
          assert pipe.safety_checker is None

          lowerCamelCase__				   =      pipe("example prompt"			,   num_inference_steps=2).images[0]
          assert image is not None

          # check that there's no error when saving a pipeline with one of the models being None
          with tempfile.TemporaryDirectory() as tmpdirname:
                 pipe.save_pretrained(UpperCamelCase)
                 lowerCamelCase__				   =      StableDiffusionPipeline.from_pretrained(UpperCamelCase)

          # sanity check that the pipeline still works
          assert pipe.safety_checker is None
          lowerCamelCase__				   =      pipe("example prompt"			,   num_inference_steps=2).images[0]
          assert image is not None







   @unittest.skipIf(torch_device != "cuda"			,   "This test requires a GPU")
   def     __UpperCAmelCase      ( self):
          lowerCamelCase__				   =      self.dummy_cond_unet
          lowerCamelCase__				   =      PNDMScheduler(skip_prk_steps=UpperCamelCase)
          lowerCamelCase__				   =      self.dummy_vae
          lowerCamelCase__				   =      self.dummy_text_encoder
          lowerCamelCase__				   =      CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")

          # put models in fp16
          lowerCamelCase__				   =      unet.half()
          lowerCamelCase__				   =      vae.half()
          lowerCamelCase__				   =      bert.half()

          # make sure here that pndm scheduler skips prk
          lowerCamelCase__				   =      StableDiffusionPipeline(
              unet=UpperCamelCase			,   scheduler=UpperCamelCase			,   vae=UpperCamelCase			,   text_encoder=UpperCamelCase			,   tokenizer=UpperCamelCase			,   safety_checker=UpperCamelCase			,   feature_extractor=self.dummy_extractor			,   )
          lowerCamelCase__				   =      sd_pipe.to(UpperCamelCase)
          sd_pipe.set_progress_bar_config(disable=UpperCamelCase)

          lowerCamelCase__				   =      "A painting of a squirrel eating a burger"
          lowerCamelCase__				   =      sd_pipe([prompt]			,   num_inference_steps=2			,   output_type="np").images

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







@nightly
@require_torch_gpu
class 		snake_case_       (						unittest.TestCase						):



   """simple docstring"""





   def     __UpperCAmelCase      ( self):
          # clean up the VRAM after each test
          super().tearDown()
          gc.collect()
          torch.cuda.empty_cache()







   def     __UpperCAmelCase      ( self):
          lowerCamelCase__				   =      "cuda"
          lowerCamelCase__				   =      StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
          lowerCamelCase__				   =      pipe.to(UpperCamelCase)
          pipe.set_progress_bar_config(disable=UpperCamelCase)

          lowerCamelCase__				   =      "a photo of a cat"
          lowerCamelCase__				   =      {
              "editing_prompt": ["sunglasses"],
              "reverse_editing_direction": [False],
              "edit_warmup_steps": 10,
              "edit_guidance_scale": 6,
              "edit_threshold": 0.9_5,
              "edit_momentum_scale": 0.5,
              "edit_mom_beta": 0.6,
          }

          lowerCamelCase__				   =      3
          lowerCamelCase__				   =      7

          # no sega enabled
          lowerCamelCase__				   =      torch.Generator(UpperCamelCase)
          generator.manual_seed(UpperCamelCase)
          lowerCamelCase__				   =      pipe(
              [prompt]			,   generator=UpperCamelCase			,   guidance_scale=UpperCamelCase			,   num_inference_steps=50			,   output_type="np"			,   width=5_12			,   height=5_12			,   )

          lowerCamelCase__				   =      output.images
          lowerCamelCase__				   =      image[0, -3:, -3:, -1]
          lowerCamelCase__				   =      [
              0.3_4_6_7_3_1_1_3,
              0.3_8_4_9_2_7_3_3,
              0.3_7_5_9_7_3_5_2,
              0.3_4_0_8_6_3_3_5,
              0.3_5_6_5_0_7_4_8,
              0.3_5_5_7_9_2_0_5,
              0.3_3_8_4_7_6_3,
              0.3_4_3_4_0_2_3_6,
              0.3_5_7_3_2_7_1,
          ]

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

          assert np.abs(image_slice.flatten() - expected_slice).max() < 1E-2

          # with sega enabled
          # generator = torch.manual_seed(seed)
          generator.manual_seed(UpperCamelCase)
          lowerCamelCase__				   =      pipe(
              [prompt]			,   generator=UpperCamelCase			,   guidance_scale=UpperCamelCase			,   num_inference_steps=50			,   output_type="np"			,   width=5_12			,   height=5_12			,   **UpperCamelCase			,   )

          lowerCamelCase__				   =      output.images
          lowerCamelCase__				   =      image[0, -3:, -3:, -1]
          lowerCamelCase__				   =      [
              0.4_1_8_8_7_8_2_6,
              0.3_7_7_2_8_7_6_6,
              0.3_0_1_3_8_2_7_2,
              0.4_1_4_1_6_3_3_5,
              0.4_1_6_6_4_9_8_5,
              0.3_6_2_8_3_3_9_2,
              0.3_6_1_9_1_2_4_6,
              0.4_3_3_6_4_4_6_5,
              0.4_3_0_0_1_7_3_2,
          ]

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

          assert np.abs(image_slice.flatten() - expected_slice).max() < 1E-2







   def     __UpperCAmelCase      ( self):
          lowerCamelCase__				   =      "cuda"
          lowerCamelCase__				   =      StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
          lowerCamelCase__				   =      pipe.to(UpperCamelCase)
          pipe.set_progress_bar_config(disable=UpperCamelCase)

          lowerCamelCase__				   =      "an image of a crowded boulevard, realistic, 4k"
          lowerCamelCase__				   =      {
              "editing_prompt": "crowd, crowded, people",
              "reverse_editing_direction": True,
              "edit_warmup_steps": 10,
              "edit_guidance_scale": 8.3,
              "edit_threshold": 0.9,
              "edit_momentum_scale": 0.5,
              "edit_mom_beta": 0.6,
          }

          lowerCamelCase__				   =      9
          lowerCamelCase__				   =      7

          # no sega enabled
          lowerCamelCase__				   =      torch.Generator(UpperCamelCase)
          generator.manual_seed(UpperCamelCase)
          lowerCamelCase__				   =      pipe(
              [prompt]			,   generator=UpperCamelCase			,   guidance_scale=UpperCamelCase			,   num_inference_steps=50			,   output_type="np"			,   width=5_12			,   height=5_12			,   )

          lowerCamelCase__				   =      output.images
          lowerCamelCase__				   =      image[0, -3:, -3:, -1]
          lowerCamelCase__				   =      [
              0.4_3_4_9_7_9_9_8,
              0.9_1_8_1_4_0_6_5,
              0.7_5_4_0_7_3_9,
              0.5_5_5_8_0_2_0_5,
              0.8_4_6_7_2_6_5,
              0.5_3_8_9_6_9_1,
              0.6_2_5_7_4_5_0_6,
              0.5_8_8_9_7_7_6_3,
              0.5_0_9_2_6_7_5_7,
          ]

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

          assert np.abs(image_slice.flatten() - expected_slice).max() < 1E-2

          # with sega enabled
          # generator = torch.manual_seed(seed)
          generator.manual_seed(UpperCamelCase)
          lowerCamelCase__				   =      pipe(
              [prompt]			,   generator=UpperCamelCase			,   guidance_scale=UpperCamelCase			,   num_inference_steps=50			,   output_type="np"			,   width=5_12			,   height=5_12			,   **UpperCamelCase			,   )

          lowerCamelCase__				   =      output.images
          lowerCamelCase__				   =      image[0, -3:, -3:, -1]
          lowerCamelCase__				   =      [
              0.3_0_8_9_7_1_9,
              0.3_0_5_0_0_1_4_4,
              0.2_9_0_1_6_0_4_2,
              0.3_0_6_3_0_9_6_4,
              0.3_2_5_6_8_7,
              0.2_9_4_1_9_2_2_5,
              0.2_9_0_8_0_9_1,
              0.2_8_7_2_3_5_9_8,
              0.2_7_6_9_6_2_9_4,
          ]

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

          assert np.abs(image_slice.flatten() - expected_slice).max() < 1E-2







   def     __UpperCAmelCase      ( self):
          lowerCamelCase__				   =      "cuda"
          lowerCamelCase__				   =      StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
          lowerCamelCase__				   =      pipe.to(UpperCamelCase)
          pipe.set_progress_bar_config(disable=UpperCamelCase)

          lowerCamelCase__				   =      "a castle next to a river"
          lowerCamelCase__				   =      {
              "editing_prompt": ["boat on a river, boat", "monet, impression, sunrise"],
              "reverse_editing_direction": False,
              "edit_warmup_steps": [15, 18],
              "edit_guidance_scale": 6,
              "edit_threshold": [0.9, 0.8],
              "edit_momentum_scale": 0.5,
              "edit_mom_beta": 0.6,
          }

          lowerCamelCase__				   =      48
          lowerCamelCase__				   =      7

          # no sega enabled
          lowerCamelCase__				   =      torch.Generator(UpperCamelCase)
          generator.manual_seed(UpperCamelCase)
          lowerCamelCase__				   =      pipe(
              [prompt]			,   generator=UpperCamelCase			,   guidance_scale=UpperCamelCase			,   num_inference_steps=50			,   output_type="np"			,   width=5_12			,   height=5_12			,   )

          lowerCamelCase__				   =      output.images
          lowerCamelCase__				   =      image[0, -3:, -3:, -1]
          lowerCamelCase__				   =      [
              0.7_5_1_6_3_5_5_5,
              0.7_6_0_3_7_1_4_5,
              0.6_1_7_8_5,
              0.9_1_8_9_6_7_3,
              0.8_6_2_7_7_0_1,
              0.8_5_1_8_9_6_9_4,
              0.8_5_1_2_8_1_3,
              0.8_7_0_1_2_0_7_6,
              0.8_3_1_2_8_5_7,
          ]

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

          assert np.abs(image_slice.flatten() - expected_slice).max() < 1E-2

          # with sega enabled
          # generator = torch.manual_seed(seed)
          generator.manual_seed(UpperCamelCase)
          lowerCamelCase__				   =      pipe(
              [prompt]			,   generator=UpperCamelCase			,   guidance_scale=UpperCamelCase			,   num_inference_steps=50			,   output_type="np"			,   width=5_12			,   height=5_12			,   **UpperCamelCase			,   )

          lowerCamelCase__				   =      output.images
          lowerCamelCase__				   =      image[0, -3:, -3:, -1]
          lowerCamelCase__				   =      [
              0.7_3_5_5_3_3_6_5,
              0.7_5_3_7_2_7_1,
              0.7_4_3_4_1_9_0_5,
              0.6_6_4_8_0_3_5_6,
              0.6_4_7_2_9_2_5,
              0.6_3_0_3_9_4_1_6,
              0.6_4_8_1_2_9_0_5,
              0.6_7_4_9_7_1_7,
              0.6_5_1_7_1_0_2,
          ]

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

          assert np.abs(image_slice.flatten() - expected_slice).max() < 1E-2







   def     __UpperCAmelCase      ( self):
          lowerCamelCase__				   =      "cuda"
          lowerCamelCase__				   =      StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5"			,   torch_dtype=torch.floataa)
          lowerCamelCase__				   =      pipe.to(UpperCamelCase)
          pipe.set_progress_bar_config(disable=UpperCamelCase)

          lowerCamelCase__				   =      "a photo of a cat"
          lowerCamelCase__				   =      {
              "editing_prompt": ["sunglasses"],
              "reverse_editing_direction": [False],
              "edit_warmup_steps": 10,
              "edit_guidance_scale": 6,
              "edit_threshold": 0.9_5,
              "edit_momentum_scale": 0.5,
              "edit_mom_beta": 0.6,
          }

          lowerCamelCase__				   =      3
          lowerCamelCase__				   =      7

          # no sega enabled
          lowerCamelCase__				   =      torch.Generator(UpperCamelCase)
          generator.manual_seed(UpperCamelCase)
          lowerCamelCase__				   =      pipe(
              [prompt]			,   generator=UpperCamelCase			,   guidance_scale=UpperCamelCase			,   num_inference_steps=50			,   output_type="np"			,   width=5_12			,   height=5_12			,   )

          lowerCamelCase__				   =      output.images
          lowerCamelCase__				   =      image[0, -3:, -3:, -1]
          lowerCamelCase__				   =      [
              0.3_4_8_8_7_6_9_5,
              0.3_8_7_6_9_5_3,
              0.3_7_5,
              0.3_4_4_2_3_8_2_8,
              0.3_5_8_1_5_4_3,
              0.3_5_7_1_7_7_7_3,
              0.3_3_8_3_7_8_9,
              0.3_4_5_7_0_3_1_2,
              0.3_5_9_3_7_5,
          ]

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

          assert np.abs(image_slice.flatten() - expected_slice).max() < 1E-2

          # with sega enabled
          # generator = torch.manual_seed(seed)
          generator.manual_seed(UpperCamelCase)
          lowerCamelCase__				   =      pipe(
              [prompt]			,   generator=UpperCamelCase			,   guidance_scale=UpperCamelCase			,   num_inference_steps=50			,   output_type="np"			,   width=5_12			,   height=5_12			,   **UpperCamelCase			,   )

          lowerCamelCase__				   =      output.images
          lowerCamelCase__				   =      image[0, -3:, -3:, -1]
          lowerCamelCase__				   =      [
              0.4_2_2_8_5_1_5_6,
              0.3_6_9_1_4_0_6_2,
              0.2_9_0_7_7_1_4_8,
              0.4_2_0_4_1_0_1_6,
              0.4_1_9_1_8_9_4_5,
              0.3_5_4_9_8_0_4_7,
              0.3_6_1_8_1_6_4,
              0.4_4_2_3_8_2_8,
              0.4_3_1_1_5_2_3_4,
          ]

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

          assert np.abs(image_slice.flatten() - expected_slice).max() < 1E-2

