





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.stable_diffusion_safe import StableDiffusionPipelineSafe as StableDiffusionPipeline
from diffusers.utils import floats_tensor, nightly, torch_device
from diffusers.utils.testing_utils import require_torch_gpu



class 							_UpperCamelCase					(	unittest.TestCase     ):

       def lowercase (		self: int		)       -> str:




              """simple docstring"""


              super().tearDown()
              gc.collect()
              torch.cuda.empty_cache()

       @property
       def lowercase (		self: Optional[int]		)       -> str:




              """simple docstring"""


              UpperCamelCase_          =  1
              UpperCamelCase_          =  3
              UpperCamelCase_          =  (32, 32)

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

       @property
       def lowercase (		self: List[str]		)       -> Dict:




              """simple docstring"""


              torch.manual_seed(0		)
              UpperCamelCase_          =  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 lowercase (		self: List[Any]		)       -> int:




              """simple docstring"""


              torch.manual_seed(0		)
              UpperCamelCase_          =  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 lowercase (		self: Tuple		)       -> Dict:




              """simple docstring"""


              torch.manual_seed(0		)
              UpperCamelCase_          =  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=1000					,   )
              return CLIPTextModel(_SCREAMING_SNAKE_CASE		)

       @property
       def lowercase (		self: int		)       -> List[Any]:




              """simple docstring"""


              def extract(*_SCREAMING_SNAKE_CASE: Dict					,   **_SCREAMING_SNAKE_CASE: Any		):



                     class 							_UpperCamelCase					:

                            def __init__(		self: List[Any]		)       -> List[Any]:




                                   """simple docstring"""


                                   UpperCamelCase_          =  torch.ones([0]		)



                            def lowercase (		self: List[Any]					,   _SCREAMING_SNAKE_CASE: Optional[int]		)       -> Dict:




                                   """simple docstring"""


                                   self.pixel_values.to(_SCREAMING_SNAKE_CASE		)
                                   return self

                     return Out()

              return extract

       def lowercase (		self: Union[str, Any]		)       -> Optional[Any]:




              """simple docstring"""


              UpperCamelCase_          =  "cpu"  # ensure determinism for the device-dependent torch.Generator
              UpperCamelCase_          =  self.dummy_cond_unet
              UpperCamelCase_          =  DDIMScheduler(
                  beta_start=0.0_00_85					,   beta_end=0.0_12					,   beta_schedule="scaled_linear"					,   clip_sample=_SCREAMING_SNAKE_CASE					,   set_alpha_to_one=_SCREAMING_SNAKE_CASE					,   )

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

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

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

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

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

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

              assert image.shape == (1, 64, 64, 3)
              UpperCamelCase_          =  np.array([0.57_56, 0.61_18, 0.50_05, 0.50_41, 0.54_71, 0.47_26, 0.49_76, 0.48_65, 0.48_64]		)

              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 lowercase (		self: Any		)       -> Optional[Any]:




              """simple docstring"""


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

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

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

              UpperCamelCase_          =  output.images

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

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

              assert image.shape == (1, 64, 64, 3)
              UpperCamelCase_          =  np.array([0.51_25, 0.57_16, 0.48_28, 0.50_60, 0.56_50, 0.47_68, 0.51_85, 0.48_95, 0.49_93]		)

              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 lowercase (		self: Optional[int]		)       -> List[Any]:




              """simple docstring"""


              UpperCamelCase_          =  StableDiffusionPipeline.from_pretrained(
                  "hf-internal-testing/tiny-stable-diffusion-lms-pipe"					,   safety_checker=_SCREAMING_SNAKE_CASE		)
              assert isinstance(_SCREAMING_SNAKE_CASE					,   _SCREAMING_SNAKE_CASE		)
              assert isinstance(pipe.scheduler					,   _SCREAMING_SNAKE_CASE		)
              assert pipe.safety_checker is None

              UpperCamelCase_          =  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(_SCREAMING_SNAKE_CASE		)
                     UpperCamelCase_          =  StableDiffusionPipeline.from_pretrained(_SCREAMING_SNAKE_CASE		)

              # sanity check that the pipeline still works
              assert pipe.safety_checker is None
              UpperCamelCase_          =  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 lowercase (		self: List[Any]		)       -> Dict:




              """simple docstring"""


              UpperCamelCase_          =  self.dummy_cond_unet
              UpperCamelCase_          =  PNDMScheduler(skip_prk_steps=_SCREAMING_SNAKE_CASE		)
              UpperCamelCase_          =  self.dummy_vae
              UpperCamelCase_          =  self.dummy_text_encoder
              UpperCamelCase_          =  CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"		)

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

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

              UpperCamelCase_          =  "A painting of a squirrel eating a burger"
              UpperCamelCase_          =  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 lowercase (		self: List[Any]		)       -> Union[str, Any]:




              """simple docstring"""


              super().tearDown()
              gc.collect()
              torch.cuda.empty_cache()

       def lowercase (		self: Any		)       -> Optional[Any]:




              """simple docstring"""


              UpperCamelCase_          =  StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5"					,   safety_checker=_SCREAMING_SNAKE_CASE		)
              UpperCamelCase_          =  LMSDiscreteScheduler.from_config(sd_pipe.scheduler.config		)
              UpperCamelCase_          =  sd_pipe.to(_SCREAMING_SNAKE_CASE		)
              sd_pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE		)

              UpperCamelCase_          =  (
                  "portrait of girl with smokey eyes makeup in abandoned hotel, grange clothes, redshift, wide high angle"
                  " coloured polaroid photograph with flash, kodak film, hyper real, stunning moody cinematography, with"
                  " anamorphic lenses, by maripol, fallen angels by wong kar - wai, style of suspiria and neon demon and"
                  " children from bahnhof zoo, detailed "
              )
              UpperCamelCase_          =  4003660346
              UpperCamelCase_          =  7

              # without safety guidance (sld_guidance_scale = 0)
              UpperCamelCase_          =  torch.manual_seed(_SCREAMING_SNAKE_CASE		)
              UpperCamelCase_          =  sd_pipe(
                  [prompt]					,   generator=_SCREAMING_SNAKE_CASE					,   guidance_scale=_SCREAMING_SNAKE_CASE					,   num_inference_steps=50					,   output_type="np"					,   width=512					,   height=512					,   sld_guidance_scale=0					,   )

              UpperCamelCase_          =  output.images
              UpperCamelCase_          =  image[0, -3:, -3:, -1]
              UpperCamelCase_          =  [0.22_78, 0.22_31, 0.22_49, 0.23_33, 0.23_03, 0.18_85, 0.22_73, 0.21_44, 0.21_76]

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

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

              # without safety guidance (strong configuration)
              UpperCamelCase_          =  torch.manual_seed(_SCREAMING_SNAKE_CASE		)
              UpperCamelCase_          =  sd_pipe(
                  [prompt]					,   generator=_SCREAMING_SNAKE_CASE					,   guidance_scale=_SCREAMING_SNAKE_CASE					,   num_inference_steps=50					,   output_type="np"					,   width=512					,   height=512					,   sld_guidance_scale=2000					,   sld_warmup_steps=7					,   sld_threshold=0.0_25					,   sld_momentum_scale=0.5					,   sld_mom_beta=0.7					,   )

              UpperCamelCase_          =  output.images
              UpperCamelCase_          =  image[0, -3:, -3:, -1]
              UpperCamelCase_          =  [0.23_83, 0.22_76, 0.2_36, 0.21_92, 0.21_86, 0.20_53, 0.19_71, 0.19_01, 0.17_19]

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

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

       def lowercase (		self: Optional[int]		)       -> str:




              """simple docstring"""


              UpperCamelCase_          =  StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5"					,   safety_checker=_SCREAMING_SNAKE_CASE		)
              UpperCamelCase_          =  LMSDiscreteScheduler.from_config(sd_pipe.scheduler.config		)
              UpperCamelCase_          =  sd_pipe.to(_SCREAMING_SNAKE_CASE		)
              sd_pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE		)

              UpperCamelCase_          =  "padme amidala taking a bath artwork, safe for work, no nudity"
              UpperCamelCase_          =  2734971755
              UpperCamelCase_          =  7

              UpperCamelCase_          =  torch.manual_seed(_SCREAMING_SNAKE_CASE		)
              UpperCamelCase_          =  sd_pipe(
                  [prompt]					,   generator=_SCREAMING_SNAKE_CASE					,   guidance_scale=_SCREAMING_SNAKE_CASE					,   num_inference_steps=50					,   output_type="np"					,   width=512					,   height=512					,   sld_guidance_scale=0					,   )

              UpperCamelCase_          =  output.images
              UpperCamelCase_          =  image[0, -3:, -3:, -1]
              UpperCamelCase_          =  [0.35_02, 0.36_22, 0.33_96, 0.36_42, 0.34_78, 0.33_18, 0.35, 0.33_48, 0.32_97]

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

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

              UpperCamelCase_          =  torch.manual_seed(_SCREAMING_SNAKE_CASE		)
              UpperCamelCase_          =  sd_pipe(
                  [prompt]					,   generator=_SCREAMING_SNAKE_CASE					,   guidance_scale=_SCREAMING_SNAKE_CASE					,   num_inference_steps=50					,   output_type="np"					,   width=512					,   height=512					,   sld_guidance_scale=2000					,   sld_warmup_steps=7					,   sld_threshold=0.0_25					,   sld_momentum_scale=0.5					,   sld_mom_beta=0.7					,   )

              UpperCamelCase_          =  output.images
              UpperCamelCase_          =  image[0, -3:, -3:, -1]
              UpperCamelCase_          =  [0.55_31, 0.52_06, 0.48_95, 0.51_56, 0.51_82, 0.47_51, 0.48_02, 0.48_03, 0.44_43]

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

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



       def lowercase (		self: List[Any]		)       -> Optional[int]:




              """simple docstring"""


              UpperCamelCase_          =  StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5"		)
              UpperCamelCase_          =  sd_pipe.to(_SCREAMING_SNAKE_CASE		)
              sd_pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE		)

              UpperCamelCase_          =  (
                  "the four horsewomen of the apocalypse, painting by tom of finland, gaston bussiere, craig mullins, j. c."
                  " leyendecker"
              )
              UpperCamelCase_          =  1044355234
              UpperCamelCase_          =  12

              UpperCamelCase_          =  torch.manual_seed(_SCREAMING_SNAKE_CASE		)
              UpperCamelCase_          =  sd_pipe(
                  [prompt]					,   generator=_SCREAMING_SNAKE_CASE					,   guidance_scale=_SCREAMING_SNAKE_CASE					,   num_inference_steps=50					,   output_type="np"					,   width=512					,   height=512					,   sld_guidance_scale=0					,   )

              UpperCamelCase_          =  output.images
              UpperCamelCase_          =  image[0, -3:, -3:, -1]
              UpperCamelCase_          =  np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]		)

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

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

              UpperCamelCase_          =  torch.manual_seed(_SCREAMING_SNAKE_CASE		)
              UpperCamelCase_          =  sd_pipe(
                  [prompt]					,   generator=_SCREAMING_SNAKE_CASE					,   guidance_scale=_SCREAMING_SNAKE_CASE					,   num_inference_steps=50					,   output_type="np"					,   width=512					,   height=512					,   sld_guidance_scale=2000					,   sld_warmup_steps=7					,   sld_threshold=0.0_25					,   sld_momentum_scale=0.5					,   sld_mom_beta=0.7					,   )

              UpperCamelCase_          =  output.images
              UpperCamelCase_          =  image[0, -3:, -3:, -1]
              UpperCamelCase_          =  np.array([0.58_18, 0.62_85, 0.68_35, 0.60_19, 0.6_25, 0.67_54, 0.60_96, 0.63_34, 0.65_61]		)
              assert image.shape == (1, 512, 512, 3)

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








_UpperCAmelCase        =				{
    'a': 'AAAAA',
    'b': 'AAAAB',
    'c': 'AAABA',
    'd': 'AAABB',
    'e': 'AABAA',
    'f': 'AABAB',
    'g': 'AABBA',
    'h': 'AABBB',
    'i': 'ABAAA',
    'j': 'BBBAA',
    'k': 'ABAAB',
    'l': 'ABABA',
    'm': 'ABABB',
    'n': 'ABBAA',
    'o': 'ABBAB',
    'p': 'ABBBA',
    'q': 'ABBBB',
    'r': 'BAAAA',
    's': 'BAAAB',
    't': 'BAABA',
    'u': 'BAABB',
    'v': 'BBBAB',
    'w': 'BABAA',
    'x': 'BABAB',
    'y': 'BABBA',
    'z': 'BABBB',
    ' ': ' ',
}


_UpperCAmelCase        =				{value: key for key, value in encode_dict.items()}




def  lowerCAmelCase_ (							UpperCamelCase_  )     ->		str:
       UpperCamelCase_          =  ""
       for letter in word.lower():
              if letter.isalpha() or letter == " ":
                     encoded += encode_dict[letter]
              else:
                     raise Exception("encode() accepts only letters of the alphabet and spaces"  )
       return encoded




def  lowerCAmelCase_ (							UpperCamelCase_  )     ->		str:
       if set(UpperCamelCase_  ) - {"A", "B", " "} != set():
              raise Exception("decode() accepts only 'A', 'B' and spaces"  )
       UpperCamelCase_          =  ""
       for word in coded.split():
              while len(UpperCamelCase_  ) != 0:
                     decoded += decode_dict[word[:5]]
                     UpperCamelCase_          =  word[5:]
              decoded += " "
       return decoded.strip()


if __name__ == "__main__":
     from doctest import testmod

     testmod()

