






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


      def _UpperCamelCase    (     self					:						str    ):





            """simple docstring"""







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


      @property
      def _UpperCamelCase    (     self					:						Any    ):





            """simple docstring"""







            __snake_case   :     str	      =1
            __snake_case   :     List[Any]	      =3
            __snake_case   :     Optional[Any]	      =(3_2, 3_2)

            __snake_case   :     str	      =floats_tensor((batch_size, num_channels) + sizes				, rng=random.Random(0    )    ).to(a_    )
            return image


      @property
      def _UpperCamelCase    (     self					:						List[Any]    ):





            """simple docstring"""







            torch.manual_seed(0    )
            __snake_case   :     Optional[Any]	      =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 _UpperCamelCase    (     self					:						str    ):





            """simple docstring"""







            torch.manual_seed(0    )
            __snake_case   :     Any	      =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 _UpperCamelCase    (     self					:						Dict    ):





            """simple docstring"""







            torch.manual_seed(0    )
            __snake_case   :     Optional[int]	      =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(a_    )


      @property
      def _UpperCamelCase    (     self					:						Optional[Any]    ):





            """simple docstring"""







            def extract(*a					:						List[Any]				, **a					:						List[str]    ):


                  class   _lowercase			:


                        def __init__(     self					:						List[Any]    ):





                              """simple docstring"""







                              __snake_case   :     Optional[int]	      =torch.ones([0]    )







                        def _UpperCamelCase    (     self					:						Any				, a					:						List[str]    ):





                              """simple docstring"""







                              self.pixel_values.to(a_    )
                              return self



                  return Out()

            return extract


      def _UpperCamelCase    (     self					:						str    ):





            """simple docstring"""







            __snake_case   :     Tuple	      ="""cpu"""  # ensure determinism for the device-dependent torch.Generator
            __snake_case   :     Optional[Any]	      =self.dummy_cond_unet
            __snake_case   :     List[str]	      =DDIMScheduler(
                beta_start=0.0_0_0_8_5				, beta_end=0.0_1_2				, beta_schedule='''scaled_linear'''				, clip_sample=a_				, set_alpha_to_one=a_				, )

            __snake_case   :     Union[str, Any]	      =self.dummy_vae
            __snake_case   :     str	      =self.dummy_text_encoder
            __snake_case   :     Any	      =CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''    )

            # make sure here that pndm scheduler skips prk
            __snake_case   :     int	      =StableDiffusionPipeline(
                unet=a_				, scheduler=a_				, vae=a_				, text_encoder=a_				, tokenizer=a_				, safety_checker=a_				, feature_extractor=self.dummy_extractor				, )
            __snake_case   :     Optional[int]	      =sd_pipe.to(a_    )
            sd_pipe.set_progress_bar_config(disable=a_    )

            __snake_case   :     List[Any]	      ="""A painting of a squirrel eating a burger"""

            __snake_case   :     Any	      =torch.Generator(device=a_    ).manual_seed(0    )
            __snake_case   :     Any	      =sd_pipe([prompt]				, generator=a_				, guidance_scale=6.0				, num_inference_steps=2				, output_type='''np'''    )
            __snake_case   :     Optional[Any]	      =output.images

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

            __snake_case   :     Union[str, Any]	      =image[0, -3:, -3:, -1]
            __snake_case   :     List[Any]	      =image_from_tuple[0, -3:, -3:, -1]

            assert image.shape == (1, 6_4, 6_4, 3)
            __snake_case   :     Dict	      =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					:						int    ):





            """simple docstring"""







            __snake_case   :     Tuple	      ="""cpu"""  # ensure determinism for the device-dependent torch.Generator
            __snake_case   :     Optional[Any]	      =self.dummy_cond_unet
            __snake_case   :     int	      =PNDMScheduler(skip_prk_steps=a_    )
            __snake_case   :     Any	      =self.dummy_vae
            __snake_case   :     List[str]	      =self.dummy_text_encoder
            __snake_case   :     List[str]	      =CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''    )

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

            __snake_case   :     int	      ="""A painting of a squirrel eating a burger"""
            __snake_case   :     Optional[int]	      =torch.Generator(device=a_    ).manual_seed(0    )
            __snake_case   :     int	      =sd_pipe([prompt]				, generator=a_				, guidance_scale=6.0				, num_inference_steps=2				, output_type='''np'''    )

            __snake_case   :     List[str]	      =output.images

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

            __snake_case   :     List[str]	      =image[0, -3:, -3:, -1]
            __snake_case   :     Optional[Any]	      =image_from_tuple[0, -3:, -3:, -1]

            assert image.shape == (1, 6_4, 6_4, 3)
            __snake_case   :     Optional[int]	      =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					:						List[str]    ):





            """simple docstring"""







            __snake_case   :     Optional[Any]	      =StableDiffusionPipeline.from_pretrained(
                '''hf-internal-testing/tiny-stable-diffusion-lms-pipe'''				, safety_checker=a_    )
            assert isinstance(a_				, a_    )
            assert isinstance(pipe.scheduler				, a_    )
            assert pipe.safety_checker is None

            __snake_case   :     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(a_    )
                  __snake_case   :     Union[str, Any]	      =StableDiffusionPipeline.from_pretrained(a_    )

            # sanity check that the pipeline still works
            assert pipe.safety_checker is None
            __snake_case   :     Tuple	      =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					:						int    ):





            """simple docstring"""







            __snake_case   :     Any	      =self.dummy_cond_unet
            __snake_case   :     Optional[Any]	      =PNDMScheduler(skip_prk_steps=a_    )
            __snake_case   :     Dict	      =self.dummy_vae
            __snake_case   :     Dict	      =self.dummy_text_encoder
            __snake_case   :     Optional[int]	      =CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''    )

            # put models in fp16
            __snake_case   :     List[Any]	      =unet.half()
            __snake_case   :     List[str]	      =vae.half()
            __snake_case   :     Dict	      =bert.half()

            # make sure here that pndm scheduler skips prk
            __snake_case   :     int	      =StableDiffusionPipeline(
                unet=a_				, scheduler=a_				, vae=a_				, text_encoder=a_				, tokenizer=a_				, safety_checker=a_				, feature_extractor=self.dummy_extractor				, )
            __snake_case   :     Optional[int]	      =sd_pipe.to(a_    )
            sd_pipe.set_progress_bar_config(disable=a_    )

            __snake_case   :     int	      ="""A painting of a squirrel eating a burger"""
            __snake_case   :     str	      =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   _lowercase			(  unittest.TestCase    ):


      def _UpperCamelCase    (     self					:						Tuple    ):





            """simple docstring"""







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


      def _UpperCamelCase    (     self					:						Any    ):





            """simple docstring"""







            __snake_case   :     int	      ="""cuda"""
            __snake_case   :     List[str]	      =StableDiffusionPipeline.from_pretrained('''runwayml/stable-diffusion-v1-5'''    )
            __snake_case   :     Optional[int]	      =pipe.to(a_    )
            pipe.set_progress_bar_config(disable=a_    )

            __snake_case   :     Any	      ="""a photo of a cat"""
            __snake_case   :     int	      ={
                """editing_prompt""": ["""sunglasses"""],
                """reverse_editing_direction""": [False],
                """edit_warmup_steps""": 1_0,
                """edit_guidance_scale""": 6,
                """edit_threshold""": 0.9_5,
                """edit_momentum_scale""": 0.5,
                """edit_mom_beta""": 0.6,
            }

            __snake_case   :     Optional[int]	      =3
            __snake_case   :     str	      =7

            # no sega enabled
            __snake_case   :     int	      =torch.Generator(a_    )
            generator.manual_seed(a_    )
            __snake_case   :     str	      =pipe(
                [prompt]				, generator=a_				, guidance_scale=a_				, num_inference_steps=5_0				, output_type='''np'''				, width=5_1_2				, height=5_1_2				, )

            __snake_case   :     str	      =output.images
            __snake_case   :     Optional[Any]	      =image[0, -3:, -3:, -1]
            __snake_case   :     Tuple	      =[
                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_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(a_    )
            __snake_case   :     Optional[Any]	      =pipe(
                [prompt]				, generator=a_				, guidance_scale=a_				, num_inference_steps=5_0				, output_type='''np'''				, width=5_1_2				, height=5_1_2				, **a_				, )

            __snake_case   :     List[str]	      =output.images
            __snake_case   :     Dict	      =image[0, -3:, -3:, -1]
            __snake_case   :     Optional[int]	      =[
                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_1_2, 5_1_2, 3)

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


      def _UpperCamelCase    (     self					:						Tuple    ):





            """simple docstring"""







            __snake_case   :     List[Any]	      ="""cuda"""
            __snake_case   :     Union[str, Any]	      =StableDiffusionPipeline.from_pretrained('''runwayml/stable-diffusion-v1-5'''    )
            __snake_case   :     Optional[Any]	      =pipe.to(a_    )
            pipe.set_progress_bar_config(disable=a_    )

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

            __snake_case   :     str	      =9
            __snake_case   :     Optional[Any]	      =7

            # no sega enabled
            __snake_case   :     Tuple	      =torch.Generator(a_    )
            generator.manual_seed(a_    )
            __snake_case   :     List[Any]	      =pipe(
                [prompt]				, generator=a_				, guidance_scale=a_				, num_inference_steps=5_0				, output_type='''np'''				, width=5_1_2				, height=5_1_2				, )

            __snake_case   :     List[Any]	      =output.images
            __snake_case   :     Any	      =image[0, -3:, -3:, -1]
            __snake_case   :     str	      =[
                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_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(a_    )
            __snake_case   :     Dict	      =pipe(
                [prompt]				, generator=a_				, guidance_scale=a_				, num_inference_steps=5_0				, output_type='''np'''				, width=5_1_2				, height=5_1_2				, **a_				, )

            __snake_case   :     int	      =output.images
            __snake_case   :     List[str]	      =image[0, -3:, -3:, -1]
            __snake_case   :     Optional[int]	      =[
                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_1_2, 5_1_2, 3)

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


      def _UpperCamelCase    (     self					:						List[Any]    ):





            """simple docstring"""







            __snake_case   :     List[str]	      ="""cuda"""
            __snake_case   :     Dict	      =StableDiffusionPipeline.from_pretrained('''runwayml/stable-diffusion-v1-5'''    )
            __snake_case   :     Tuple	      =pipe.to(a_    )
            pipe.set_progress_bar_config(disable=a_    )

            __snake_case   :     List[Any]	      ="""a castle next to a river"""
            __snake_case   :     Tuple	      ={
                """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,
            }

            __snake_case   :     List[Any]	      =4_8
            __snake_case   :     str	      =7

            # no sega enabled
            __snake_case   :     str	      =torch.Generator(a_    )
            generator.manual_seed(a_    )
            __snake_case   :     Dict	      =pipe(
                [prompt]				, generator=a_				, guidance_scale=a_				, num_inference_steps=5_0				, output_type='''np'''				, width=5_1_2				, height=5_1_2				, )

            __snake_case   :     Optional[Any]	      =output.images
            __snake_case   :     Union[str, Any]	      =image[0, -3:, -3:, -1]
            __snake_case   :     List[Any]	      =[
                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_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(a_    )
            __snake_case   :     Dict	      =pipe(
                [prompt]				, generator=a_				, guidance_scale=a_				, num_inference_steps=5_0				, output_type='''np'''				, width=5_1_2				, height=5_1_2				, **a_				, )

            __snake_case   :     Dict	      =output.images
            __snake_case   :     Optional[int]	      =image[0, -3:, -3:, -1]
            __snake_case   :     List[str]	      =[
                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_1_2, 5_1_2, 3)

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







      def _UpperCamelCase    (     self					:						Any    ):





            """simple docstring"""







            __snake_case   :     Union[str, Any]	      ="""cuda"""
            __snake_case   :     List[str]	      =StableDiffusionPipeline.from_pretrained('''runwayml/stable-diffusion-v1-5'''				, torch_dtype=torch.floataa    )
            __snake_case   :     List[str]	      =pipe.to(a_    )
            pipe.set_progress_bar_config(disable=a_    )

            __snake_case   :     Union[str, Any]	      ="""a photo of a cat"""
            __snake_case   :     Tuple	      ={
                """editing_prompt""": ["""sunglasses"""],
                """reverse_editing_direction""": [False],
                """edit_warmup_steps""": 1_0,
                """edit_guidance_scale""": 6,
                """edit_threshold""": 0.9_5,
                """edit_momentum_scale""": 0.5,
                """edit_mom_beta""": 0.6,
            }

            __snake_case   :     Any	      =3
            __snake_case   :     Dict	      =7

            # no sega enabled
            __snake_case   :     Optional[int]	      =torch.Generator(a_    )
            generator.manual_seed(a_    )
            __snake_case   :     Any	      =pipe(
                [prompt]				, generator=a_				, guidance_scale=a_				, num_inference_steps=5_0				, output_type='''np'''				, width=5_1_2				, height=5_1_2				, )

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

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

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







'''simple docstring'''






def        snake_case_   (	lowerCAmelCase_			,  lowerCAmelCase_						)->						float:





       '''simple docstring'''


       return a + b


if __name__ == "__main__":
   A_				:					Tuple							=      5
   A_				:					Tuple							=      6
   print(f"""The sum of {a} + {b} is {add(a, b)}""")



