


'''simple docstring'''




import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    StableDiffusionAttendAndExcitePipeline,
    UNetaDConditionModel,
)
from diffusers.utils import load_numpy, skip_mps, slow
from diffusers.utils.testing_utils import require_torch_gpu

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


__snake_case    						=False

@skip_mps
class 		UpperCAmelCase_      (  __lowercase       ,	__lowercase       ,	__lowercase       ,	unittest.TestCase		):
     lowerCamelCase				:	List[Any]		=     StableDiffusionAttendAndExcitePipeline
     lowerCamelCase				:	Any		=     False
     lowerCamelCase				:	List[str]		=     TEXT_TO_IMAGE_PARAMS
     lowerCamelCase				:	Tuple		=     TEXT_TO_IMAGE_BATCH_PARAMS.union({'''token_indices'''}		)
     lowerCamelCase				:	Dict		=     TEXT_TO_IMAGE_IMAGE_PARAMS
     lowerCamelCase				:	Optional[Any]		=     TEXT_TO_IMAGE_IMAGE_PARAMS





     @classmethod
     def 							__UpperCAmelCase      ( cls       :      str    )				->  Union[str, Any]:
        super().setUpClass()
        torch.use_deterministic_algorithms(UpperCAmelCase__    )





     @classmethod
     def 							__UpperCAmelCase      ( cls       :      int    )				->  List[str]:
        super().tearDownClass()
        torch.use_deterministic_algorithms(UpperCAmelCase__    )





     def 							__UpperCAmelCase      ( self       :      int    )				->  int:
        torch.manual_seed(0    )
        lowerCAmelCase     =			UNetaDConditionModel(
            block_out_channels=(3_2, 6_4)    ,      layers_per_block=1    ,      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    ,      attention_head_dim=(2, 4)    ,      use_linear_projection=UpperCAmelCase__    ,      )
        lowerCAmelCase     =			DDIMScheduler(
            beta_start=0.00_085    ,      beta_end=0.012    ,      beta_schedule='scaled_linear'    ,      clip_sample=UpperCAmelCase__    ,      set_alpha_to_one=UpperCAmelCase__    ,      )
        torch.manual_seed(0    )
        lowerCAmelCase     =			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    ,      sample_size=1_2_8    ,      )
        torch.manual_seed(0    )
        lowerCAmelCase     =			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    ,      hidden_act='gelu'    ,      projection_dim=5_1_2    ,      )
        lowerCAmelCase     =			CLIPTextModel(UpperCAmelCase__    )
        lowerCAmelCase     =			CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'    )

        lowerCAmelCase     =			{
            'unet': unet,
            'scheduler': scheduler,
            'vae': vae,
            'text_encoder': text_encoder,
            'tokenizer': tokenizer,
            'safety_checker': None,
            'feature_extractor': None,
        }

        return components





     def 							__UpperCAmelCase      ( self       :      Any    ,      UpperCAmelCase__       :      int    ,      UpperCAmelCase__       :      Union[str, Any]=0    )				->  Tuple:
        if str(UpperCAmelCase__    ).startswith('mps'    ):
           lowerCAmelCase     =			torch.manual_seed(UpperCAmelCase__    )
        else:
           lowerCAmelCase     =			torch.Generator(device=UpperCAmelCase__    ).manual_seed(UpperCAmelCase__    )
        lowerCAmelCase     =			lowerCAmelCase     =			{
            'prompt': 'a cat and a frog',
            'token_indices': [2, 5],
            'generator': generator,
            'num_inference_steps': 1,
            'guidance_scale': 6.0,
            'output_type': 'numpy',
            'max_iter_to_alter': 2,
            'thresholds': {0: 0.7},
        }
        return inputs





     def 							__UpperCAmelCase      ( self       :      Optional[int]    )				->  List[Any]:
        lowerCAmelCase     =			'cpu'

        lowerCAmelCase     =			self.get_dummy_components()
        lowerCAmelCase     =			self.pipeline_class(**UpperCAmelCase__    )
        pipe.to(UpperCAmelCase__    )
        pipe.set_progress_bar_config(disable=UpperCAmelCase__    )

        lowerCAmelCase     =			self.get_dummy_inputs(UpperCAmelCase__    )
        lowerCAmelCase     =			pipe(**UpperCAmelCase__    ).images
        lowerCAmelCase     =			image[0, -3:, -3:, -1]

        self.assertEqual(image.shape    ,      (1, 6_4, 6_4, 3)    )
        lowerCAmelCase     =			np.array(
            [0.63_905_364, 0.62_897_307, 0.48_599_017, 0.5_133_624, 0.5_550_048, 0.45_769_516, 0.50_326_973, 0.5_023_139, 0.45_384_496]    )
        lowerCAmelCase     =			np.abs(image_slice.flatten() - expected_slice    ).max()
        self.assertLessEqual(UpperCAmelCase__    ,      1E-3    )





     def 							__UpperCAmelCase      ( self       :      List[Any]    )				->  Optional[int]:
        super().test_cpu_offload_forward_pass(expected_max_diff=5E-4    )





     def 							__UpperCAmelCase      ( self       :      str    )				->  int:
        # NOTE: Larger batch sizes cause this test to timeout, only test on smaller batches
        self._test_inference_batch_consistent(batch_sizes=[1, 2]    )





     def 							__UpperCAmelCase      ( self       :      Optional[int]    )				->  Optional[int]:
        self._test_inference_batch_single_identical(batch_size=2    ,      expected_max_diff=7E-4    )





     def 							__UpperCAmelCase      ( self       :      List[Any]    )				->  int:
        super().test_dict_tuple_outputs_equivalent(expected_max_difference=3E-3    )





     def 							__UpperCAmelCase      ( self       :      Union[str, Any]    )				->  Tuple:
        super().test_pt_np_pil_outputs_equivalent(expected_max_diff=5E-4    )





     def 							__UpperCAmelCase      ( self       :      Optional[int]    )				->  Any:
        super().test_save_load_local(expected_max_difference=5E-4    )





     def 							__UpperCAmelCase      ( self       :      List[Any]    )				->  List[str]:
        super().test_save_load_optional_components(expected_max_difference=4E-4    )




@require_torch_gpu
@slow
class 		UpperCAmelCase_      (  unittest.TestCase		):





     @classmethod
     def 							__UpperCAmelCase      ( cls       :      List[Any]    )				->  List[Any]:
        super().setUpClass()
        torch.use_deterministic_algorithms(UpperCAmelCase__    )





     @classmethod
     def 							__UpperCAmelCase      ( cls       :      str    )				->  Optional[Any]:
        super().tearDownClass()
        torch.use_deterministic_algorithms(UpperCAmelCase__    )





     def 							__UpperCAmelCase      ( self       :      Tuple    )				->  int:
        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()





     def 							__UpperCAmelCase      ( self       :      List[str]    )				->  Optional[Any]:
        lowerCAmelCase     =			torch.manual_seed(5_1    )

        lowerCAmelCase     =			StableDiffusionAttendAndExcitePipeline.from_pretrained(
            'CompVis/stable-diffusion-v1-4'    ,      safety_checker=UpperCAmelCase__    ,      torch_dtype=torch.floataa    )
        pipe.to('cuda'    )

        lowerCAmelCase     =			'a painting of an elephant with glasses'
        lowerCAmelCase     =			[5, 7]

        lowerCAmelCase     =			pipe(
            prompt=UpperCAmelCase__    ,      token_indices=UpperCAmelCase__    ,      guidance_scale=7.5    ,      generator=UpperCAmelCase__    ,      num_inference_steps=5    ,      max_iter_to_alter=5    ,      output_type='numpy'    ,      ).images[0]

        lowerCAmelCase     =			load_numpy(
            'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/attend-and-excite/elephant_glasses.npy'    )
        assert np.abs((expected_image - image).max()    ) < 5E-1









'''simple docstring'''




import gc
import unittest

import numpy as np
import torch

from diffusers import StableDiffusionKDiffusionPipeline
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()

@slow
@require_torch_gpu
class 		UpperCAmelCase_      (  unittest.TestCase		):





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





     def 							__UpperCAmelCase      ( self       :      Tuple    )				->  Any:
        lowerCAmelCase     =			StableDiffusionKDiffusionPipeline.from_pretrained('CompVis/stable-diffusion-v1-4'    )
        lowerCAmelCase     =			sd_pipe.to(UpperCAmelCase__    )
        sd_pipe.set_progress_bar_config(disable=UpperCAmelCase__    )

        sd_pipe.set_scheduler('sample_euler'    )

        lowerCAmelCase     =			'A painting of a squirrel eating a burger'
        lowerCAmelCase     =			torch.manual_seed(0    )
        lowerCAmelCase     =			sd_pipe([prompt]    ,      generator=UpperCAmelCase__    ,      guidance_scale=9.0    ,      num_inference_steps=2_0    ,      output_type='np'    )

        lowerCAmelCase     =			output.images

        lowerCAmelCase     =			image[0, -3:, -3:, -1]

        assert image.shape == (1, 5_1_2, 5_1_2, 3)
        lowerCAmelCase     =			np.array([0.0_447, 0.0_492, 0.0_468, 0.0_408, 0.0_383, 0.0_408, 0.0_354, 0.0_380, 0.0_339]    )

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





     def 							__UpperCAmelCase      ( self       :      List[str]    )				->  Dict:
        lowerCAmelCase     =			StableDiffusionKDiffusionPipeline.from_pretrained('stabilityai/stable-diffusion-2-1-base'    )
        lowerCAmelCase     =			sd_pipe.to(UpperCAmelCase__    )
        sd_pipe.set_progress_bar_config(disable=UpperCAmelCase__    )

        sd_pipe.set_scheduler('sample_euler'    )

        lowerCAmelCase     =			'A painting of a squirrel eating a burger'
        lowerCAmelCase     =			torch.manual_seed(0    )
        lowerCAmelCase     =			sd_pipe([prompt]    ,      generator=UpperCAmelCase__    ,      guidance_scale=9.0    ,      num_inference_steps=2_0    ,      output_type='np'    )

        lowerCAmelCase     =			output.images

        lowerCAmelCase     =			image[0, -3:, -3:, -1]

        assert image.shape == (1, 5_1_2, 5_1_2, 3)
        lowerCAmelCase     =			np.array([0.1_237, 0.1_320, 0.1_438, 0.1_359, 0.1_390, 0.1_132, 0.1_277, 0.1_175, 0.1_112]    )

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





     def 							__UpperCAmelCase      ( self       :      Optional[Any]    )				->  List[str]:
        lowerCAmelCase     =			StableDiffusionKDiffusionPipeline.from_pretrained('stabilityai/stable-diffusion-2-1-base'    )
        lowerCAmelCase     =			sd_pipe.to(UpperCAmelCase__    )
        sd_pipe.set_progress_bar_config(disable=UpperCAmelCase__    )

        sd_pipe.set_scheduler('sample_dpmpp_2m'    )

        lowerCAmelCase     =			'A painting of a squirrel eating a burger'
        lowerCAmelCase     =			torch.manual_seed(0    )
        lowerCAmelCase     =			sd_pipe(
            [prompt]    ,      generator=UpperCAmelCase__    ,      guidance_scale=7.5    ,      num_inference_steps=1_5    ,      output_type='np'    ,      use_karras_sigmas=UpperCAmelCase__    ,      )

        lowerCAmelCase     =			output.images

        lowerCAmelCase     =			image[0, -3:, -3:, -1]

        assert image.shape == (1, 5_1_2, 5_1_2, 3)
        lowerCAmelCase     =			np.array(
            [0.11_381_689, 0.12_112_921, 0.1_389_457, 0.12_549_606, 0.1_244_964, 0.10_831_517, 0.11_562_866, 0.10_867_816, 0.10_499_048]    )

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





