


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





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





     @property
     def 					__A (self	)					->       Any:
      _lowercase											=1
      _lowercase											=3
      _lowercase											=(3_2, 3_2)

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





     @property
     def 					__A (self	)					->       Union[str, Any]:
      torch.manual_seed(0	)
      _lowercase											=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	)
      _lowercase											=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	)					->       int:
      torch.manual_seed(0	)
      _lowercase											=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(UpperCAmelCase	)





     @property
     def 					__A (self	)					->       Tuple:
      def extract(*UpperCAmelCase		,					**UpperCAmelCase	):







       class 				lowerCamelCase__       :





            def __init__(self	)					->       Dict:
             _lowercase											=torch.ones([0]	)





            def 					__A (self		,					UpperCAmelCase	)					->       Any:
             self.pixel_values.to(UpperCAmelCase	)
             return self



       return Out()

      return extract





     def 					__A (self	)					->       List[str]:
      _lowercase											='''cpu'''  # ensure determinism for the device-dependent torch.Generator
      _lowercase											=self.dummy_cond_unet
      _lowercase											=DDIMScheduler(
          beta_start=0.0_0085		,					beta_end=0.012		,					beta_schedule='''scaled_linear'''		,					clip_sample=UpperCAmelCase		,					set_alpha_to_one=UpperCAmelCase		,					)

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

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

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

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

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

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

      assert image.shape == (1, 6_4, 6_4, 3)
      _lowercase											=np.array([0.5753, 0.6114, 0.5001, 0.5034, 0.5470, 0.4729, 0.4971, 0.4867, 0.4867]	)

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

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

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

      _lowercase											=output.images

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

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

      assert image.shape == (1, 6_4, 6_4, 3)
      _lowercase											=np.array([0.5122, 0.5712, 0.4825, 0.5053, 0.5646, 0.4769, 0.5179, 0.4894, 0.4994]	)

      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	)					->       Union[str, Any]:
      _lowercase											=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

      _lowercase											=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	)
       _lowercase											=StableDiffusionPipeline.from_pretrained(UpperCAmelCase	)

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

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

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

      _lowercase											='''A painting of a squirrel eating a burger'''
      _lowercase											=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 				lowerCamelCase__       (	unittest.TestCase):





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





     def 					__A (self	)					->       str:
      _lowercase											='''cuda'''
      _lowercase											=StableDiffusionPipeline.from_pretrained('''runwayml/stable-diffusion-v1-5'''	)
      _lowercase											=pipe.to(UpperCAmelCase	)
      pipe.set_progress_bar_config(disable=UpperCAmelCase	)

      _lowercase											='''a photo of a cat'''
      _lowercase											={
          '''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,
      }

      _lowercase											=3
      _lowercase											=7

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

      _lowercase											=output.images
      _lowercase											=image[0, -3:, -3:, -1]
      _lowercase											=[
          0.3467_3113,
          0.3849_2733,
          0.3759_7352,
          0.3408_6335,
          0.3565_0748,
          0.3557_9205,
          0.338_4763,
          0.3434_0236,
          0.357_3271,
      ]

      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(UpperCAmelCase	)
      _lowercase											=pipe(
          [prompt]		,					generator=UpperCAmelCase		,					guidance_scale=UpperCAmelCase		,					num_inference_steps=5_0		,					output_type='''np'''		,					width=5_1_2		,					height=5_1_2		,					**UpperCAmelCase		,					)

      _lowercase											=output.images
      _lowercase											=image[0, -3:, -3:, -1]
      _lowercase											=[
          0.4188_7826,
          0.3772_8766,
          0.3013_8272,
          0.4141_6335,
          0.4166_4985,
          0.3628_3392,
          0.3619_1246,
          0.4336_4465,
          0.4300_1732,
      ]

      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]:
      _lowercase											='''cuda'''
      _lowercase											=StableDiffusionPipeline.from_pretrained('''runwayml/stable-diffusion-v1-5'''	)
      _lowercase											=pipe.to(UpperCAmelCase	)
      pipe.set_progress_bar_config(disable=UpperCAmelCase	)

      _lowercase											='''an image of a crowded boulevard, realistic, 4k'''
      _lowercase											={
          '''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,
      }

      _lowercase											=9
      _lowercase											=7

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

      _lowercase											=output.images
      _lowercase											=image[0, -3:, -3:, -1]
      _lowercase											=[
          0.4349_7998,
          0.9181_4065,
          0.754_0739,
          0.5558_0205,
          0.846_7265,
          0.538_9691,
          0.6257_4506,
          0.5889_7763,
          0.5092_6757,
      ]

      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(UpperCAmelCase	)
      _lowercase											=pipe(
          [prompt]		,					generator=UpperCAmelCase		,					guidance_scale=UpperCAmelCase		,					num_inference_steps=5_0		,					output_type='''np'''		,					width=5_1_2		,					height=5_1_2		,					**UpperCAmelCase		,					)

      _lowercase											=output.images
      _lowercase											=image[0, -3:, -3:, -1]
      _lowercase											=[
          0.308_9719,
          0.3050_0144,
          0.2901_6042,
          0.3063_0964,
          0.32_5687,
          0.2941_9225,
          0.290_8091,
          0.2872_3598,
          0.2769_6294,
      ]

      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	)					->       str:
      _lowercase											='''cuda'''
      _lowercase											=StableDiffusionPipeline.from_pretrained('''runwayml/stable-diffusion-v1-5'''	)
      _lowercase											=pipe.to(UpperCAmelCase	)
      pipe.set_progress_bar_config(disable=UpperCAmelCase	)

      _lowercase											='''a castle next to a river'''
      _lowercase											={
          '''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,
      }

      _lowercase											=4_8
      _lowercase											=7

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

      _lowercase											=output.images
      _lowercase											=image[0, -3:, -3:, -1]
      _lowercase											=[
          0.7516_3555,
          0.7603_7145,
          0.6_1785,
          0.918_9673,
          0.862_7701,
          0.8518_9694,
          0.851_2813,
          0.8701_2076,
          0.831_2857,
      ]

      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(UpperCAmelCase	)
      _lowercase											=pipe(
          [prompt]		,					generator=UpperCAmelCase		,					guidance_scale=UpperCAmelCase		,					num_inference_steps=5_0		,					output_type='''np'''		,					width=5_1_2		,					height=5_1_2		,					**UpperCAmelCase		,					)

      _lowercase											=output.images
      _lowercase											=image[0, -3:, -3:, -1]
      _lowercase											=[
          0.7355_3365,
          0.753_7271,
          0.7434_1905,
          0.6648_0356,
          0.647_2925,
          0.6303_9416,
          0.6481_2905,
          0.674_9717,
          0.651_7102,
      ]

      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	)					->       Tuple:
      _lowercase											='''cuda'''
      _lowercase											=StableDiffusionPipeline.from_pretrained('''runwayml/stable-diffusion-v1-5'''		,					torch_dtype=torch.floataa	)
      _lowercase											=pipe.to(UpperCAmelCase	)
      pipe.set_progress_bar_config(disable=UpperCAmelCase	)

      _lowercase											='''a photo of a cat'''
      _lowercase											={
          '''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,
      }

      _lowercase											=3
      _lowercase											=7

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

      _lowercase											=output.images
      _lowercase											=image[0, -3:, -3:, -1]
      _lowercase											=[
          0.3488_7695,
          0.387_6953,
          0.375,
          0.3442_3828,
          0.358_1543,
          0.3571_7773,
          0.338_3789,
          0.3457_0312,
          0.35_9375,
      ]

      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(UpperCAmelCase	)
      _lowercase											=pipe(
          [prompt]		,					generator=UpperCAmelCase		,					guidance_scale=UpperCAmelCase		,					num_inference_steps=5_0		,					output_type='''np'''		,					width=5_1_2		,					height=5_1_2		,					**UpperCAmelCase		,					)

      _lowercase											=output.images
      _lowercase											=image[0, -3:, -3:, -1]
      _lowercase											=[
          0.4228_5156,
          0.3691_4062,
          0.2907_7148,
          0.4204_1016,
          0.4191_8945,
          0.3549_8047,
          0.361_8164,
          0.442_3828,
          0.4311_5234,
      ]

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

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





from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


UpperCAmelCase__            =							{'''configuration_visual_bert''': ['''VISUAL_BERT_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''VisualBertConfig''']}

try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  UpperCAmelCase__            =							[
      '''VISUAL_BERT_PRETRAINED_MODEL_ARCHIVE_LIST''',
      '''VisualBertForMultipleChoice''',
      '''VisualBertForPreTraining''',
      '''VisualBertForQuestionAnswering''',
      '''VisualBertForRegionToPhraseAlignment''',
      '''VisualBertForVisualReasoning''',
      '''VisualBertLayer''',
      '''VisualBertModel''',
      '''VisualBertPreTrainedModel''',
  ]


if TYPE_CHECKING:
  from .configuration_visual_bert import VISUAL_BERT_PRETRAINED_CONFIG_ARCHIVE_MAP, VisualBertConfig

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_visual_bert import (
        VISUAL_BERT_PRETRAINED_MODEL_ARCHIVE_LIST,
        VisualBertForMultipleChoice,
        VisualBertForPreTraining,
        VisualBertForQuestionAnswering,
        VisualBertForRegionToPhraseAlignment,
        VisualBertForVisualReasoning,
        VisualBertLayer,
        VisualBertModel,
        VisualBertPreTrainedModel,
    )


else:
  import sys

  UpperCAmelCase__            =							_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)

