



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





     """simple docstring"""






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





     @property
     def 						A__  (					self :							Any ):
          A__				=   1
          A__				=   3
          A__				=   (3_2, 3_2)

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





     @property
     def 						A__  (					self :							int ):
          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	,      )
          return model





     @property
     def 						A__  (					self :							Any ):
          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	,      )
          return model





     @property
     def 						A__  (					self :							Dict ):
          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	,      )
          return CLIPTextModel(_lowerCamelCase )





     @property
     def 						A__  (					self :							Optional[int] ):
          def extract(*_lowerCamelCase :							str	,      **_lowerCamelCase :							Dict ):

               class       UpperCamelCase  :





                    """simple docstring"""






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





                    def 						A__  (					self :							int	,      _lowerCamelCase :							List[Any] ):
                         self.pixel_values.to(_lowerCamelCase )
                         return self


               return Out()

          return extract





     def 						A__  (					self :							str ):
          A__				=   '''cpu'''  # ensure determinism for the device-dependent torch.Generator
          A__				=   self.dummy_cond_unet
          A__				=   DDIMScheduler(
              beta_start=0.00_085	,      beta_end=0.012	,      beta_schedule='''scaled_linear'''	,      clip_sample=_lowerCamelCase	,      set_alpha_to_one=_lowerCamelCase	,      )

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

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

          A__				=   '''A painting of a squirrel eating a burger'''

          A__				=   torch.Generator(device=_lowerCamelCase ).manual_seed(0 )
          A__				=   sd_pipe([prompt]	,      generator=_lowerCamelCase	,      guidance_scale=6.0	,      num_inference_steps=2	,      output_type='''np''' )
          A__				=   output.images

          A__				=   torch.Generator(device=_lowerCamelCase ).manual_seed(0 )
          A__				=   sd_pipe(
              [prompt]	,      generator=_lowerCamelCase	,      guidance_scale=6.0	,      num_inference_steps=2	,      output_type='''np'''	,      return_dict=_lowerCamelCase	,      )[0]

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

          assert image.shape == (1, 6_4, 6_4, 3)
          A__				=   np.array([0.5_753, 0.6_114, 0.5_001, 0.5_034, 0.5_470, 0.4_729, 0.4_971, 0.4_867, 0.4_867] )

          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 						A__  (					self :							Dict ):
          A__				=   '''cpu'''  # ensure determinism for the device-dependent torch.Generator
          A__				=   self.dummy_cond_unet
          A__				=   PNDMScheduler(skip_prk_steps=_lowerCamelCase )
          A__				=   self.dummy_vae
          A__				=   self.dummy_text_encoder
          A__				=   CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip''' )

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

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

          A__				=   output.images

          A__				=   torch.Generator(device=_lowerCamelCase ).manual_seed(0 )
          A__				=   sd_pipe(
              [prompt]	,      generator=_lowerCamelCase	,      guidance_scale=6.0	,      num_inference_steps=2	,      output_type='''np'''	,      return_dict=_lowerCamelCase	,      )[0]

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

          assert image.shape == (1, 6_4, 6_4, 3)
          A__				=   np.array([0.5_122, 0.5_712, 0.4_825, 0.5_053, 0.5_646, 0.4_769, 0.5_179, 0.4_894, 0.4_994] )

          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 						A__  (					self :							str ):
          A__				=   StableDiffusionPipeline.from_pretrained(
              '''hf-internal-testing/tiny-stable-diffusion-lms-pipe'''	,      safety_checker=_lowerCamelCase )
          assert isinstance(_lowerCamelCase	,      _lowerCamelCase )
          assert isinstance(pipe.scheduler	,      _lowerCamelCase )
          assert pipe.safety_checker is None

          A__				=   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(_lowerCamelCase )
               A__				=   StableDiffusionPipeline.from_pretrained(_lowerCamelCase )

          # sanity check that the pipeline still works
          assert pipe.safety_checker is None
          A__				=   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 						A__  (					self :							str ):
          A__				=   self.dummy_cond_unet
          A__				=   PNDMScheduler(skip_prk_steps=_lowerCamelCase )
          A__				=   self.dummy_vae
          A__				=   self.dummy_text_encoder
          A__				=   CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip''' )

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

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

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

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


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





     """simple docstring"""






     def 						A__  (					self :							List[Any] ):
          # clean up the VRAM after each test
          super().tearDown()
          gc.collect()
          torch.cuda.empty_cache()





     def 						A__  (					self :							Any ):
          A__				=   '''cuda'''
          A__				=   StableDiffusionPipeline.from_pretrained('''runwayml/stable-diffusion-v1-5''' )
          A__				=   pipe.to(_lowerCamelCase )
          pipe.set_progress_bar_config(disable=_lowerCamelCase )

          A__				=   '''a photo of a cat'''
          A__				=   {
              '''editing_prompt''': ['''sunglasses'''],
              '''reverse_editing_direction''': [False],
              '''edit_warmup_steps''': 1_0,
              '''edit_guidance_scale''': 6,
              '''edit_threshold''': 0.95,
              '''edit_momentum_scale''': 0.5,
              '''edit_mom_beta''': 0.6,
          }

          A__				=   3
          A__				=   7

          # no sega enabled
          A__				=   torch.Generator(_lowerCamelCase )
          generator.manual_seed(_lowerCamelCase )
          A__				=   pipe(
              [prompt]	,      generator=_lowerCamelCase	,      guidance_scale=_lowerCamelCase	,      num_inference_steps=5_0	,      output_type='''np'''	,      width=5_1_2	,      height=5_1_2	,      )

          A__				=   output.images
          A__				=   image[0, -3:, -3:, -1]
          A__				=   [
              0.34_673_113,
              0.38_492_733,
              0.37_597_352,
              0.34_086_335,
              0.35_650_748,
              0.35_579_205,
              0.3_384_763,
              0.34_340_236,
              0.3_573_271,
          ]

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

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

          # with sega enabled
          # generator = torch.manual_seed(seed)
          generator.manual_seed(_lowerCamelCase )
          A__				=   pipe(
              [prompt]	,      generator=_lowerCamelCase	,      guidance_scale=_lowerCamelCase	,      num_inference_steps=5_0	,      output_type='''np'''	,      width=5_1_2	,      height=5_1_2	,      **_lowerCamelCase	,      )

          A__				=   output.images
          A__				=   image[0, -3:, -3:, -1]
          A__				=   [
              0.41_887_826,
              0.37_728_766,
              0.30_138_272,
              0.41_416_335,
              0.41_664_985,
              0.36_283_392,
              0.36_191_246,
              0.43_364_465,
              0.43_001_732,
          ]

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

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





     def 						A__  (					self :							Union[str, Any] ):
          A__				=   '''cuda'''
          A__				=   StableDiffusionPipeline.from_pretrained('''runwayml/stable-diffusion-v1-5''' )
          A__				=   pipe.to(_lowerCamelCase )
          pipe.set_progress_bar_config(disable=_lowerCamelCase )

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

          A__				=   9
          A__				=   7

          # no sega enabled
          A__				=   torch.Generator(_lowerCamelCase )
          generator.manual_seed(_lowerCamelCase )
          A__				=   pipe(
              [prompt]	,      generator=_lowerCamelCase	,      guidance_scale=_lowerCamelCase	,      num_inference_steps=5_0	,      output_type='''np'''	,      width=5_1_2	,      height=5_1_2	,      )

          A__				=   output.images
          A__				=   image[0, -3:, -3:, -1]
          A__				=   [
              0.43_497_998,
              0.91_814_065,
              0.7_540_739,
              0.55_580_205,
              0.8_467_265,
              0.5_389_691,
              0.62_574_506,
              0.58_897_763,
              0.50_926_757,
          ]

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

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

          # with sega enabled
          # generator = torch.manual_seed(seed)
          generator.manual_seed(_lowerCamelCase )
          A__				=   pipe(
              [prompt]	,      generator=_lowerCamelCase	,      guidance_scale=_lowerCamelCase	,      num_inference_steps=5_0	,      output_type='''np'''	,      width=5_1_2	,      height=5_1_2	,      **_lowerCamelCase	,      )

          A__				=   output.images
          A__				=   image[0, -3:, -3:, -1]
          A__				=   [
              0.3_089_719,
              0.30_500_144,
              0.29_016_042,
              0.30_630_964,
              0.325_687,
              0.29_419_225,
              0.2_908_091,
              0.28_723_598,
              0.27_696_294,
          ]

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

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





     def 						A__  (					self :							Optional[Any] ):
          A__				=   '''cuda'''
          A__				=   StableDiffusionPipeline.from_pretrained('''runwayml/stable-diffusion-v1-5''' )
          A__				=   pipe.to(_lowerCamelCase )
          pipe.set_progress_bar_config(disable=_lowerCamelCase )

          A__				=   '''a castle next to a river'''
          A__				=   {
              '''editing_prompt''': ['''boat on a river, boat''', '''monet, impression, sunrise'''],
              '''reverse_editing_direction''': False,
              '''edit_warmup_steps''': [1_5, 1_8],
              '''edit_guidance_scale''': 6,
              '''edit_threshold''': [0.9, 0.8],
              '''edit_momentum_scale''': 0.5,
              '''edit_mom_beta''': 0.6,
          }

          A__				=   4_8
          A__				=   7

          # no sega enabled
          A__				=   torch.Generator(_lowerCamelCase )
          generator.manual_seed(_lowerCamelCase )
          A__				=   pipe(
              [prompt]	,      generator=_lowerCamelCase	,      guidance_scale=_lowerCamelCase	,      num_inference_steps=5_0	,      output_type='''np'''	,      width=5_1_2	,      height=5_1_2	,      )

          A__				=   output.images
          A__				=   image[0, -3:, -3:, -1]
          A__				=   [
              0.75_163_555,
              0.76_037_145,
              0.61_785,
              0.9_189_673,
              0.8_627_701,
              0.85_189_694,
              0.8_512_813,
              0.87_012_076,
              0.8_312_857,
          ]

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

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

          # with sega enabled
          # generator = torch.manual_seed(seed)
          generator.manual_seed(_lowerCamelCase )
          A__				=   pipe(
              [prompt]	,      generator=_lowerCamelCase	,      guidance_scale=_lowerCamelCase	,      num_inference_steps=5_0	,      output_type='''np'''	,      width=5_1_2	,      height=5_1_2	,      **_lowerCamelCase	,      )

          A__				=   output.images
          A__				=   image[0, -3:, -3:, -1]
          A__				=   [
              0.73_553_365,
              0.7_537_271,
              0.74_341_905,
              0.66_480_356,
              0.6_472_925,
              0.63_039_416,
              0.64_812_905,
              0.6_749_717,
              0.6_517_102,
          ]

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

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





     def 						A__  (					self :							int ):
          A__				=   '''cuda'''
          A__				=   StableDiffusionPipeline.from_pretrained('''runwayml/stable-diffusion-v1-5'''	,      torch_dtype=torch.floataa )
          A__				=   pipe.to(_lowerCamelCase )
          pipe.set_progress_bar_config(disable=_lowerCamelCase )

          A__				=   '''a photo of a cat'''
          A__				=   {
              '''editing_prompt''': ['''sunglasses'''],
              '''reverse_editing_direction''': [False],
              '''edit_warmup_steps''': 1_0,
              '''edit_guidance_scale''': 6,
              '''edit_threshold''': 0.95,
              '''edit_momentum_scale''': 0.5,
              '''edit_mom_beta''': 0.6,
          }

          A__				=   3
          A__				=   7

          # no sega enabled
          A__				=   torch.Generator(_lowerCamelCase )
          generator.manual_seed(_lowerCamelCase )
          A__				=   pipe(
              [prompt]	,      generator=_lowerCamelCase	,      guidance_scale=_lowerCamelCase	,      num_inference_steps=5_0	,      output_type='''np'''	,      width=5_1_2	,      height=5_1_2	,      )

          A__				=   output.images
          A__				=   image[0, -3:, -3:, -1]
          A__				=   [
              0.34_887_695,
              0.3_876_953,
              0.375,
              0.34_423_828,
              0.3_581_543,
              0.35_717_773,
              0.3_383_789,
              0.34_570_312,
              0.359_375,
          ]

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

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

          # with sega enabled
          # generator = torch.manual_seed(seed)
          generator.manual_seed(_lowerCamelCase )
          A__				=   pipe(
              [prompt]	,      generator=_lowerCamelCase	,      guidance_scale=_lowerCamelCase	,      num_inference_steps=5_0	,      output_type='''np'''	,      width=5_1_2	,      height=5_1_2	,      **_lowerCamelCase	,      )

          A__				=   output.images
          A__				=   image[0, -3:, -3:, -1]
          A__				=   [
              0.42_285_156,
              0.36_914_062,
              0.29_077_148,
              0.42_041_016,
              0.41_918_945,
              0.35_498_047,
              0.3_618_164,
              0.4_423_828,
              0.43_115_234,
          ]

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

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








"""simple docstring"""







import argparse
import logging
import sys
from unittest.mock import patch

import run_glue_with_pabee

from transformers.testing_utils import TestCasePlus


logging.basicConfig(level=logging.DEBUG)

__snake_case	:				List[str]  					=			logging.getLogger()



def   a_	(			):
     A__				=   argparse.ArgumentParser()
     parser.add_argument('''-f'''      )
     A__				=   parser.parse_args()
     return args.f

class       UpperCamelCase  (					a				):





     """simple docstring"""






     def 						A__  (					self :							Optional[int] ):
          A__				=   logging.StreamHandler(sys.stdout )
          logger.addHandler(_lowerCamelCase )

          A__				=   self.get_auto_remove_tmp_dir()
          A__				=   F'''
            run_glue_with_pabee.py
            --model_type albert
            --model_name_or_path albert-base-v2
            --data_dir ./tests/fixtures/tests_samples/MRPC/
            --output_dir {tmp_dir}
            --overwrite_output_dir
            --task_name mrpc
            --do_train
            --do_eval
            --per_gpu_train_batch_size=2
            --per_gpu_eval_batch_size=1
            --learning_rate=2e-5
            --max_steps=50
            --warmup_steps=2
            --seed=42
            --max_seq_length=128
            '''.split()

          with patch.object(_lowerCamelCase	,      '''argv'''	,      _lowerCamelCase ):
               A__				=   run_glue_with_pabee.main()
               for value in result.values():
                    self.assertGreaterEqual(_lowerCamelCase	,      0.75 )



