
'''simple docstring'''







import argparse
import json
import logging
import os
import sys
from unittest import skip
from unittest.mock import patch

import tensorflow as tf

from transformers.testing_utils import TestCasePlus, get_gpu_count, slow


lowerCAmelCase_      =   [
    os.path.join(os.path.dirname(__file__), dirname)
    for dirname in [
        "text-generation",
        "text-classification",
        "token-classification",
        "language-modeling",
        "multiple-choice",
        "question-answering",
        "summarization",
        "translation",
        "image-classification",
    ]
]
sys.path.extend(SRC_DIRS)


if SRC_DIRS is not None:
    import run_clm
    import run_image_classification
    import run_mlm
    import run_ner
    import run_qa as run_squad
    import run_summarization
    import run_swag
    import run_text_classification
    import run_translation


logging.basicConfig(level=logging.DEBUG)

lowerCAmelCase_      =   logging.getLogger()







def   lowerCAmelCase(  ):







       '''simple docstring'''






       lowerCamelCase__				   =      argparse.ArgumentParser()
       parser.add_argument("-f"      )
       lowerCamelCase__				   =      parser.parse_args()
       return args.f







def   lowerCAmelCase(  a__    :  Optional[Any]      ):







       '''simple docstring'''






       lowerCamelCase__				   =      {}
       lowerCamelCase__				   =      os.path.join(a__      , "all_results.json"      )
       if os.path.exists(a__      ):
              with open(a__      , "r"      ) as f:
                     lowerCamelCase__				   =      json.load(a__      )
       else:
              raise ValueError(f"""can't find {path}"""      )
       return results







def   lowerCAmelCase(  ):







       '''simple docstring'''






       return bool(tf.config.list_physical_devices("GPU"      )      )


lowerCAmelCase_      =   logging.StreamHandler(sys.stdout)
logger.addHandler(stream_handler)







class 		snake_case_       (						A__						):



   """simple docstring"""





   @skip("Skipping until shape inference for to_tf_dataset PR is merged.")
   def     __UpperCAmelCase      ( self):
          lowerCamelCase__				   =      self.get_auto_remove_tmp_dir()
          lowerCamelCase__				   =      f"""
            run_text_classification.py
            --model_name_or_path distilbert-base-uncased
            --output_dir {tmp_dir}
            --overwrite_output_dir
            --train_file ./tests/fixtures/tests_samples/MRPC/train.csv
            --validation_file ./tests/fixtures/tests_samples/MRPC/dev.csv
            --do_train
            --do_eval
            --per_device_train_batch_size=2
            --per_device_eval_batch_size=1
            --learning_rate=1e-4
            --max_steps=10
            --warmup_steps=2
            --seed=42
            --max_seq_length=128
            """.split()

          if is_cuda_available():
                 testargs.append("--fp16")

          with patch.object(UpperCamelCase			,   "argv"			,   UpperCamelCase):
                 run_text_classification.main()
                 # Reset the mixed precision policy so we don't break other tests
                 tf.keras.mixed_precision.set_global_policy("float32")
                 lowerCamelCase__				   =      get_results(UpperCamelCase)
                 self.assertGreaterEqual(result["eval_accuracy"]			,   0.7_5)







   def     __UpperCAmelCase      ( self):
          lowerCamelCase__				   =      self.get_auto_remove_tmp_dir()
          lowerCamelCase__				   =      f"""
            run_clm.py
            --model_name_or_path distilgpt2
            --train_file ./tests/fixtures/sample_text.txt
            --validation_file ./tests/fixtures/sample_text.txt
            --do_train
            --do_eval
            --block_size 128
            --per_device_train_batch_size 2
            --per_device_eval_batch_size 1
            --num_train_epochs 2
            --output_dir {tmp_dir}
            --overwrite_output_dir
            """.split()

          if len(tf.config.list_physical_devices("GPU")) > 1:
                 # Skipping because there are not enough batches to train the model + would need a drop_last to work.
                 return

          with patch.object(UpperCamelCase			,   "argv"			,   UpperCamelCase):
                 run_clm.main()
                 lowerCamelCase__				   =      get_results(UpperCamelCase)
                 self.assertLess(result["eval_perplexity"]			,   1_00)







   def     __UpperCAmelCase      ( self):
          lowerCamelCase__				   =      self.get_auto_remove_tmp_dir()
          lowerCamelCase__				   =      f"""
            run_mlm.py
            --model_name_or_path distilroberta-base
            --train_file ./tests/fixtures/sample_text.txt
            --validation_file ./tests/fixtures/sample_text.txt
            --max_seq_length 64
            --output_dir {tmp_dir}
            --overwrite_output_dir
            --do_train
            --do_eval
            --prediction_loss_only
            --num_train_epochs=1
            --learning_rate=1e-4
        """.split()

          with patch.object(UpperCamelCase			,   "argv"			,   UpperCamelCase):
                 run_mlm.main()
                 lowerCamelCase__				   =      get_results(UpperCamelCase)
                 self.assertLess(result["eval_perplexity"]			,   42)







   def     __UpperCAmelCase      ( self):
          # with so little data distributed training needs more epochs to get the score on par with 0/1 gpu
          lowerCamelCase__				   =      7 if get_gpu_count() > 1 else 2

          lowerCamelCase__				   =      self.get_auto_remove_tmp_dir()
          lowerCamelCase__				   =      f"""
            run_ner.py
            --model_name_or_path bert-base-uncased
            --train_file tests/fixtures/tests_samples/conll/sample.json
            --validation_file tests/fixtures/tests_samples/conll/sample.json
            --output_dir {tmp_dir}
            --overwrite_output_dir
            --do_train
            --do_eval
            --warmup_steps=2
            --learning_rate=2e-4
            --per_device_train_batch_size=2
            --per_device_eval_batch_size=2
            --num_train_epochs={epochs}
            --seed 7
        """.split()

          with patch.object(UpperCamelCase			,   "argv"			,   UpperCamelCase):
                 run_ner.main()
                 lowerCamelCase__				   =      get_results(UpperCamelCase)
                 self.assertGreaterEqual(result["accuracy"]			,   0.7_5)







   def     __UpperCAmelCase      ( self):
          lowerCamelCase__				   =      self.get_auto_remove_tmp_dir()
          lowerCamelCase__				   =      f"""
            run_qa.py
            --model_name_or_path bert-base-uncased
            --version_2_with_negative
            --train_file tests/fixtures/tests_samples/SQUAD/sample.json
            --validation_file tests/fixtures/tests_samples/SQUAD/sample.json
            --output_dir {tmp_dir}
            --overwrite_output_dir
            --max_steps=10
            --warmup_steps=2
            --do_train
            --do_eval
            --learning_rate=2e-4
            --per_device_train_batch_size=2
            --per_device_eval_batch_size=1
        """.split()

          with patch.object(UpperCamelCase			,   "argv"			,   UpperCamelCase):
                 run_squad.main()
                 lowerCamelCase__				   =      get_results(UpperCamelCase)
                 self.assertGreaterEqual(result["f1"]			,   30)
                 self.assertGreaterEqual(result["exact"]			,   30)







   def     __UpperCAmelCase      ( self):
          lowerCamelCase__				   =      self.get_auto_remove_tmp_dir()
          lowerCamelCase__				   =      f"""
            run_swag.py
            --model_name_or_path bert-base-uncased
            --train_file tests/fixtures/tests_samples/swag/sample.json
            --validation_file tests/fixtures/tests_samples/swag/sample.json
            --output_dir {tmp_dir}
            --overwrite_output_dir
            --max_steps=20
            --warmup_steps=2
            --do_train
            --do_eval
            --learning_rate=2e-4
            --per_device_train_batch_size=2
            --per_device_eval_batch_size=1
        """.split()

          with patch.object(UpperCamelCase			,   "argv"			,   UpperCamelCase):
                 run_swag.main()
                 lowerCamelCase__				   =      get_results(UpperCamelCase)
                 self.assertGreaterEqual(result["val_accuracy"]			,   0.8)







   @slow
   def     __UpperCAmelCase      ( self):
          lowerCamelCase__				   =      self.get_auto_remove_tmp_dir()
          lowerCamelCase__				   =      f"""
            run_summarization.py
            --model_name_or_path t5-small
            --train_file tests/fixtures/tests_samples/xsum/sample.json
            --validation_file tests/fixtures/tests_samples/xsum/sample.json
            --output_dir {tmp_dir}
            --overwrite_output_dir
            --max_steps=50
            --warmup_steps=8
            --do_train
            --do_eval
            --learning_rate=2e-4
            --per_device_train_batch_size=2
            --per_device_eval_batch_size=1
        """.split()

          with patch.object(UpperCamelCase			,   "argv"			,   UpperCamelCase):
                 run_summarization.main()
                 lowerCamelCase__				   =      get_results(UpperCamelCase)
                 self.assertGreaterEqual(result["rouge1"]			,   10)
                 self.assertGreaterEqual(result["rouge2"]			,   2)
                 self.assertGreaterEqual(result["rougeL"]			,   7)
                 self.assertGreaterEqual(result["rougeLsum"]			,   7)







   @slow
   def     __UpperCAmelCase      ( self):
          lowerCamelCase__				   =      self.get_auto_remove_tmp_dir()
          lowerCamelCase__				   =      f"""
            run_translation.py
            --model_name_or_path Rocketknight1/student_marian_en_ro_6_1
            --source_lang en
            --target_lang ro
            --train_file tests/fixtures/tests_samples/wmt16/sample.json
            --validation_file tests/fixtures/tests_samples/wmt16/sample.json
            --output_dir {tmp_dir}
            --overwrite_output_dir
            --warmup_steps=8
            --do_train
            --do_eval
            --learning_rate=3e-3
            --num_train_epochs 12
            --per_device_train_batch_size=2
            --per_device_eval_batch_size=1
            --source_lang en_XX
            --target_lang ro_RO
        """.split()

          with patch.object(UpperCamelCase			,   "argv"			,   UpperCamelCase):
                 run_translation.main()
                 lowerCamelCase__				   =      get_results(UpperCamelCase)
                 self.assertGreaterEqual(result["bleu"]			,   30)







   def     __UpperCAmelCase      ( self):
          lowerCamelCase__				   =      self.get_auto_remove_tmp_dir()
          lowerCamelCase__				   =      f"""
            run_image_classification.py
            --dataset_name hf-internal-testing/cats_vs_dogs_sample
            --model_name_or_path microsoft/resnet-18
            --do_train
            --do_eval
            --learning_rate 1e-4
            --per_device_train_batch_size 2
            --per_device_eval_batch_size 1
            --output_dir {tmp_dir}
            --overwrite_output_dir
            --dataloader_num_workers 16
            --num_train_epochs 2
            --train_val_split 0.1
            --seed 42
            --ignore_mismatched_sizes True
            """.split()

          with patch.object(UpperCamelCase			,   "argv"			,   UpperCamelCase):
                 run_image_classification.main()
                 lowerCamelCase__				   =      get_results(UpperCamelCase)
                 self.assertGreaterEqual(result["accuracy"]			,   0.7)



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

