




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






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






   @property
   def  _a						(	self   :       int     )->			Union[str, Any]:
     _lowercase: List[str]						     =    1
     _lowercase: List[Any]						     =    3
     _lowercase: Union[str, Any]						     =    (32, 32)

     _lowercase: Optional[int]						     =    floats_tensor((batch_size, num_channels) + sizes						,					rng=random.Random(0     )     ).to(UpperCamelCase_     )
     return image






   @property
   def  _a						(	self   :       Any     )->			List[Any]:
     torch.manual_seed(0     )
     _lowercase: Tuple						     =    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  _a						(	self   :       Any     )->			List[str]:
     torch.manual_seed(0     )
     _lowercase: Optional[Any]						     =    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  _a						(	self   :       str     )->			Dict:
     torch.manual_seed(0     )
     _lowercase: Optional[Any]						     =    CLIPTextConfig(
         bos_token_id=0						,					eos_token_id=2						,					hidden_size=32						,					intermediate_size=37						,					layer_norm_eps=1E-05						,					num_attention_heads=4						,					num_hidden_layers=5						,					pad_token_id=1						,					vocab_size=10_00						,					)
     return CLIPTextModel(UpperCamelCase_     )






   @property
   def  _a						(	self   :       Optional[Any]     )->			str:
     def extract(*UpperCamelCase_   :       str						,					**UpperCamelCase_   :       List[str]     ):

       class      UpperCAmelCase						:






          def __init__(	self   :       Optional[int]     )->			Tuple:
            _lowercase: Any						     =    torch.ones([0]     )






          def  _a						(	self   :       Tuple						,					UpperCamelCase_   :       Tuple     )->			Optional[Any]:
            self.pixel_values.to(UpperCamelCase_     )
            return self
       return Out()

     return extract






   def  _a						(	self   :       Any     )->			int:
     _lowercase: Any						     =    'cpu'  # ensure determinism for the device-dependent torch.Generator
     _lowercase: Any						     =    self.dummy_cond_unet
     _lowercase: Tuple						     =    DDIMScheduler(
         beta_start=0.00_085						,					beta_end=0.012						,					beta_schedule='scaled_linear'						,					clip_sample=UpperCamelCase_						,					set_alpha_to_one=UpperCamelCase_						,					)

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

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

     _lowercase: Union[str, Any]						     =    'A painting of a squirrel eating a burger'

     _lowercase: Tuple						     =    torch.Generator(device=UpperCamelCase_     ).manual_seed(0     )
     _lowercase: str						     =    sd_pipe([prompt]						,					generator=UpperCamelCase_						,					guidance_scale=6.0						,					num_inference_steps=2						,					output_type='np'     )
     _lowercase: Union[str, Any]						     =    output.images

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

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

     assert image.shape == (1, 64, 64, 3)
     _lowercase: Optional[int]						     =    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   :       Optional[Any]     )->			Dict:
     _lowercase: Tuple						     =    'cpu'  # ensure determinism for the device-dependent torch.Generator
     _lowercase: Union[str, Any]						     =    self.dummy_cond_unet
     _lowercase: List[Any]						     =    PNDMScheduler(skip_prk_steps=UpperCamelCase_     )
     _lowercase: List[Any]						     =    self.dummy_vae
     _lowercase: Dict						     =    self.dummy_text_encoder
     _lowercase: Union[str, Any]						     =    CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'     )

     # make sure here that pndm scheduler skips prk
     _lowercase: Any						     =    StableDiffusionPipeline(
         unet=UpperCamelCase_						,					scheduler=UpperCamelCase_						,					vae=UpperCamelCase_						,					text_encoder=UpperCamelCase_						,					tokenizer=UpperCamelCase_						,					safety_checker=UpperCamelCase_						,					feature_extractor=self.dummy_extractor						,					)
     _lowercase: List[str]						     =    sd_pipe.to(UpperCamelCase_     )
     sd_pipe.set_progress_bar_config(disable=UpperCamelCase_     )

     _lowercase: int						     =    'A painting of a squirrel eating a burger'
     _lowercase: Union[str, Any]						     =    torch.Generator(device=UpperCamelCase_     ).manual_seed(0     )
     _lowercase: List[str]						     =    sd_pipe([prompt]						,					generator=UpperCamelCase_						,					guidance_scale=6.0						,					num_inference_steps=2						,					output_type='np'     )

     _lowercase: int						     =    output.images

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

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

     assert image.shape == (1, 64, 64, 3)
     _lowercase: List[str]						     =    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   :       Tuple     )->			Dict:
     _lowercase: str						     =    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: Optional[Any]						     =    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: Optional[int]						     =    StableDiffusionPipeline.from_pretrained(UpperCamelCase_     )

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

     # put models in fp16
     _lowercase: int						     =    unet.half()
     _lowercase: Dict						     =    vae.half()
     _lowercase: Optional[Any]						     =    bert.half()

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

     _lowercase: List[str]						     =    'A painting of a squirrel eating a burger'
     _lowercase: Union[str, Any]						     =    sd_pipe([prompt]						,					num_inference_steps=2						,					output_type='np'     ).images

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

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






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






   def  _a						(	self   :       Union[str, Any]     )->			Dict:
     _lowercase: Optional[Any]						     =    'cuda'
     _lowercase: List[str]						     =    StableDiffusionPipeline.from_pretrained('runwayml/stable-diffusion-v1-5'     )
     _lowercase: int						     =    pipe.to(UpperCamelCase_     )
     pipe.set_progress_bar_config(disable=UpperCamelCase_     )

     _lowercase: Dict						     =    'a photo of a cat'
     _lowercase: Optional[Any]						     =    {
         'editing_prompt': ['sunglasses'],
         'reverse_editing_direction': [False],
         'edit_warmup_steps': 10,
         'edit_guidance_scale': 6,
         'edit_threshold': 0.95,
         'edit_momentum_scale': 0.5,
         'edit_mom_beta': 0.6,
     }

     _lowercase: Optional[Any]						     =    3
     _lowercase: int						     =    7

     # no sega enabled
     _lowercase: int						     =    torch.Generator(UpperCamelCase_     )
     generator.manual_seed(UpperCamelCase_     )
     _lowercase: Tuple						     =    pipe(
         [prompt]						,					generator=UpperCamelCase_						,					guidance_scale=UpperCamelCase_						,					num_inference_steps=50						,					output_type='np'						,					width=5_12						,					height=5_12						,					)

     _lowercase: int						     =    output.images
     _lowercase: List[Any]						     =    image[0, -3:, -3:, -1]
     _lowercase: List[str]						     =    [
         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_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_     )
     _lowercase: Tuple						     =    pipe(
         [prompt]						,					generator=UpperCamelCase_						,					guidance_scale=UpperCamelCase_						,					num_inference_steps=50						,					output_type='np'						,					width=5_12						,					height=5_12						,					**UpperCamelCase_						,					)

     _lowercase: List[str]						     =    output.images
     _lowercase: Tuple						     =    image[0, -3:, -3:, -1]
     _lowercase: Union[str, Any]						     =    [
         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_12, 5_12, 3)

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






   def  _a						(	self   :       List[Any]     )->			List[str]:
     _lowercase: Tuple						     =    'cuda'
     _lowercase: Any						     =    StableDiffusionPipeline.from_pretrained('runwayml/stable-diffusion-v1-5'     )
     _lowercase: str						     =    pipe.to(UpperCamelCase_     )
     pipe.set_progress_bar_config(disable=UpperCamelCase_     )

     _lowercase: Any						     =    'an image of a crowded boulevard, realistic, 4k'
     _lowercase: Optional[int]						     =    {
         '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,
     }

     _lowercase: Optional[Any]						     =    9
     _lowercase: int						     =    7

     # no sega enabled
     _lowercase: Union[str, Any]						     =    torch.Generator(UpperCamelCase_     )
     generator.manual_seed(UpperCamelCase_     )
     _lowercase: Optional[Any]						     =    pipe(
         [prompt]						,					generator=UpperCamelCase_						,					guidance_scale=UpperCamelCase_						,					num_inference_steps=50						,					output_type='np'						,					width=5_12						,					height=5_12						,					)

     _lowercase: Dict						     =    output.images
     _lowercase: Union[str, Any]						     =    image[0, -3:, -3:, -1]
     _lowercase: Any						     =    [
         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_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_     )
     _lowercase: int						     =    pipe(
         [prompt]						,					generator=UpperCamelCase_						,					guidance_scale=UpperCamelCase_						,					num_inference_steps=50						,					output_type='np'						,					width=5_12						,					height=5_12						,					**UpperCamelCase_						,					)

     _lowercase: Union[str, Any]						     =    output.images
     _lowercase: List[Any]						     =    image[0, -3:, -3:, -1]
     _lowercase: Tuple						     =    [
         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_12, 5_12, 3)

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






   def  _a						(	self   :       Dict     )->			Union[str, Any]:
     _lowercase: Dict						     =    'cuda'
     _lowercase: Tuple						     =    StableDiffusionPipeline.from_pretrained('runwayml/stable-diffusion-v1-5'     )
     _lowercase: Optional[Any]						     =    pipe.to(UpperCamelCase_     )
     pipe.set_progress_bar_config(disable=UpperCamelCase_     )

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

     _lowercase: int						     =    48
     _lowercase: Tuple						     =    7

     # no sega enabled
     _lowercase: List[Any]						     =    torch.Generator(UpperCamelCase_     )
     generator.manual_seed(UpperCamelCase_     )
     _lowercase: Dict						     =    pipe(
         [prompt]						,					generator=UpperCamelCase_						,					guidance_scale=UpperCamelCase_						,					num_inference_steps=50						,					output_type='np'						,					width=5_12						,					height=5_12						,					)

     _lowercase: str						     =    output.images
     _lowercase: Any						     =    image[0, -3:, -3:, -1]
     _lowercase: Union[str, Any]						     =    [
         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_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_     )
     _lowercase: List[str]						     =    pipe(
         [prompt]						,					generator=UpperCamelCase_						,					guidance_scale=UpperCamelCase_						,					num_inference_steps=50						,					output_type='np'						,					width=5_12						,					height=5_12						,					**UpperCamelCase_						,					)

     _lowercase: List[Any]						     =    output.images
     _lowercase: Tuple						     =    image[0, -3:, -3:, -1]
     _lowercase: Any						     =    [
         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_12, 5_12, 3)

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






   def  _a						(	self   :       str     )->			Optional[int]:
     _lowercase: int						     =    'cuda'
     _lowercase: int						     =    StableDiffusionPipeline.from_pretrained('runwayml/stable-diffusion-v1-5'						,					torch_dtype=torch.floataa     )
     _lowercase: Union[str, Any]						     =    pipe.to(UpperCamelCase_     )
     pipe.set_progress_bar_config(disable=UpperCamelCase_     )

     _lowercase: List[str]						     =    'a photo of a cat'
     _lowercase: Tuple						     =    {
         'editing_prompt': ['sunglasses'],
         'reverse_editing_direction': [False],
         'edit_warmup_steps': 10,
         'edit_guidance_scale': 6,
         'edit_threshold': 0.95,
         'edit_momentum_scale': 0.5,
         'edit_mom_beta': 0.6,
     }

     _lowercase: Any						     =    3
     _lowercase: Optional[int]						     =    7

     # no sega enabled
     _lowercase: int						     =    torch.Generator(UpperCamelCase_     )
     generator.manual_seed(UpperCamelCase_     )
     _lowercase: Optional[Any]						     =    pipe(
         [prompt]						,					generator=UpperCamelCase_						,					guidance_scale=UpperCamelCase_						,					num_inference_steps=50						,					output_type='np'						,					width=5_12						,					height=5_12						,					)

     _lowercase: Tuple						     =    output.images
     _lowercase: Tuple						     =    image[0, -3:, -3:, -1]
     _lowercase: str						     =    [
         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_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_     )
     _lowercase: List[str]						     =    pipe(
         [prompt]						,					generator=UpperCamelCase_						,					guidance_scale=UpperCamelCase_						,					num_inference_steps=50						,					output_type='np'						,					width=5_12						,					height=5_12						,					**UpperCamelCase_						,					)

     _lowercase: Tuple						     =    output.images
     _lowercase: Union[str, Any]						     =    image[0, -3:, -3:, -1]
     _lowercase: Optional[int]						     =    [
         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_12, 5_12, 3)

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






from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_sentencepiece_available, is_torch_available


UpperCAmelCase__						:       Union[str, Any]							=    {'configuration_bert_generation': ['BertGenerationConfig']}

try:
   if not is_sentencepiece_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   UpperCAmelCase__						:       str							=    ['BertGenerationTokenizer']

try:
   if not is_torch_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   UpperCAmelCase__						:       Tuple							=    [
       'BertGenerationDecoder',
       'BertGenerationEncoder',
       'BertGenerationPreTrainedModel',
       'load_tf_weights_in_bert_generation',
   ]


if TYPE_CHECKING:
   from .configuration_bert_generation import BertGenerationConfig

   try:
      if not is_sentencepiece_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .tokenization_bert_generation import BertGenerationTokenizer

   try:
      if not is_torch_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_bert_generation import (
          BertGenerationDecoder,
          BertGenerationEncoder,
          BertGenerationPreTrainedModel,
          load_tf_weights_in_bert_generation,
      )

else:
   import sys

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