"""simple docstring"""




from dataclasses import dataclass
from enum import Enum
from typing import List, Optional, Union

import numpy as np
import PIL
from PIL import Image

from ...utils import BaseOutput, is_torch_available, is_transformers_available

@dataclass
class        __SCREAMING_SNAKE_CASE     (   _lowerCAmelCase							):
				UpperCAmelCase	=      {
				    '''sld_warmup_steps''': 15,
				    '''sld_guidance_scale''': 20,
				    '''sld_threshold''': 0.0,
				    '''sld_momentum_scale''': 0.0,
				    '''sld_mom_beta''': 0.0,
				}
				UpperCAmelCase	=      {
				    '''sld_warmup_steps''': 10,
				    '''sld_guidance_scale''': 1_000,
				    '''sld_threshold''': 0.0_1,
				    '''sld_momentum_scale''': 0.3,
				    '''sld_mom_beta''': 0.4,
				}
				UpperCAmelCase	=      {
				    '''sld_warmup_steps''': 7,
				    '''sld_guidance_scale''': 2_000,
				    '''sld_threshold''': 0.0_2_5,
				    '''sld_momentum_scale''': 0.5,
				    '''sld_mom_beta''': 0.7,
				}
				UpperCAmelCase	=      {
				    '''sld_warmup_steps''': 0,
				    '''sld_guidance_scale''': 5_000,
				    '''sld_threshold''': 1.0,
				    '''sld_momentum_scale''': 0.5,
				    '''sld_mom_beta''': 0.7,
				}





@dataclass
class        __SCREAMING_SNAKE_CASE     (   _lowerCAmelCase							):
				UpperCAmelCase	=      42
				UpperCAmelCase	=      42
				UpperCAmelCase	=      42
				UpperCAmelCase	=      42


if is_transformers_available() and is_torch_available():
					from .pipeline_stable_diffusion_safe import StableDiffusionPipelineSafe
					from .safety_checker import SafeStableDiffusionSafetyChecker







"""simple docstring"""




import gc
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMParallelScheduler,
    DDPMParallelScheduler,
    StableDiffusionParadigmsPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import (
    enable_full_determinism,
    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 PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()

class        __SCREAMING_SNAKE_CASE     (   _lowerCAmelCase						,       _lowerCAmelCase						,       unittest.TestCase							):
				UpperCAmelCase	=      StableDiffusionParadigmsPipeline
				UpperCAmelCase	=      TEXT_TO_IMAGE_PARAMS
				UpperCAmelCase	=      TEXT_TO_IMAGE_BATCH_PARAMS
				UpperCAmelCase	=      TEXT_TO_IMAGE_IMAGE_PARAMS
				UpperCAmelCase	=      TEXT_TO_IMAGE_IMAGE_PARAMS







				def   lowercase_		(    self    :List[Any]  )      ->					Tuple:




									"""simple docstring"""



									torch.manual_seed(0  )
									lowerCamelCase__		:   int  				=					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      ,attention_head_dim=(2, 4)      ,use_linear_projection=__UpperCAmelCase      ,)
									lowerCamelCase__		:   List[Any]  				=					DDIMParallelScheduler(
									    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__		:   Dict  				=					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      ,sample_size=1_28      ,)
									torch.manual_seed(0  )
									lowerCamelCase__		:   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      ,hidden_act='''gelu'''      ,projection_dim=5_12      ,)
									lowerCamelCase__		:   Tuple  				=					CLIPTextModel(__UpperCAmelCase  )
									lowerCamelCase__		:   int  				=					CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''  )

									lowerCamelCase__		:   List[Any]  				=					{
									    '''unet''': unet,
									    '''scheduler''': scheduler,
									    '''vae''': vae,
									    '''text_encoder''': text_encoder,
									    '''tokenizer''': tokenizer,
									    '''safety_checker''': None,
									    '''feature_extractor''': None,
									}
									return components







				def   lowercase_		(    self    :Tuple      ,__UpperCAmelCase    :int      ,__UpperCAmelCase    :Dict=0  )      ->					Optional[int]:




									"""simple docstring"""



									if str(__UpperCAmelCase  ).startswith('''mps'''  ):
														lowerCamelCase__		:   Dict  				=					torch.manual_seed(__UpperCAmelCase  )
									else:
														lowerCamelCase__		:   List[str]  				=					torch.Generator(device=__UpperCAmelCase  ).manual_seed(__UpperCAmelCase  )
									lowerCamelCase__		:   Optional[int]  				=					{
									    '''prompt''': '''a photograph of an astronaut riding a horse''',
									    '''generator''': generator,
									    '''num_inference_steps''': 10,
									    '''guidance_scale''': 6.0,
									    '''output_type''': '''numpy''',
									    '''parallel''': 3,
									    '''debug''': True,
									}
									return inputs







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




									"""simple docstring"""



									lowerCamelCase__		:   Any  				=					'''cpu'''  # ensure determinism for the device-dependent torch.Generator
									lowerCamelCase__		:   int  				=					self.get_dummy_components()
									lowerCamelCase__		:   int  				=					StableDiffusionParadigmsPipeline(**__UpperCAmelCase  )
									lowerCamelCase__		:   Union[str, Any]  				=					sd_pipe.to(__UpperCAmelCase  )
									sd_pipe.set_progress_bar_config(disable=__UpperCAmelCase  )

									lowerCamelCase__		:   Dict  				=					self.get_dummy_inputs(__UpperCAmelCase  )
									lowerCamelCase__		:   Tuple  				=					sd_pipe(**__UpperCAmelCase  ).images
									lowerCamelCase__		:   List[Any]  				=					image[0, -3:, -3:, -1]
									assert image.shape == (1, 64, 64, 3)

									lowerCamelCase__		:   Tuple  				=					np.array([0.4_773, 0.5_417, 0.4_723, 0.4_925, 0.5_631, 0.4_752, 0.5_240, 0.4_935, 0.5_023]  )

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







				def   lowercase_		(    self    :int  )      ->					str:




									"""simple docstring"""



									lowerCamelCase__		:   Tuple  				=					'''cpu'''  # ensure determinism for the device-dependent torch.Generator
									lowerCamelCase__		:   Optional[int]  				=					self.get_dummy_components()
									torch.manual_seed(0  )
									lowerCamelCase__		:   Any  				=					DDPMParallelScheduler()
									torch.manual_seed(0  )
									lowerCamelCase__		:   str  				=					StableDiffusionParadigmsPipeline(**__UpperCAmelCase  )
									lowerCamelCase__		:   Tuple  				=					sd_pipe.to(__UpperCAmelCase  )
									sd_pipe.set_progress_bar_config(disable=__UpperCAmelCase  )

									lowerCamelCase__		:   int  				=					self.get_dummy_inputs(__UpperCAmelCase  )
									lowerCamelCase__		:   Optional[int]  				=					sd_pipe(**__UpperCAmelCase  ).images
									lowerCamelCase__		:   int  				=					image[0, -3:, -3:, -1]
									assert image.shape == (1, 64, 64, 3)

									lowerCamelCase__		:   Optional[int]  				=					np.array([0.3_573, 0.4_420, 0.4_960, 0.4_799, 0.3_796, 0.3_879, 0.4_819, 0.4_365, 0.4_468]  )

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







				def   lowercase_		(    self    :List[str]  )      ->					Dict:




									"""simple docstring"""



									super().test_inference_batch_consistent(batch_sizes=[1, 2]  )







				def   lowercase_		(    self    :Any  )      ->					str:




									"""simple docstring"""



									super().test_inference_batch_single_identical(batch_size=2      ,expected_max_diff=3E-3  )







				def   lowercase_		(    self    :int  )      ->					List[str]:




									"""simple docstring"""



									lowerCamelCase__		:   Optional[Any]  				=					'''cpu'''  # ensure determinism for the device-dependent torch.Generator
									lowerCamelCase__		:   str  				=					self.get_dummy_components()
									lowerCamelCase__		:   int  				=					StableDiffusionParadigmsPipeline(**__UpperCAmelCase  )
									lowerCamelCase__		:   List[Any]  				=					sd_pipe.to(__UpperCAmelCase  )
									sd_pipe.set_progress_bar_config(disable=__UpperCAmelCase  )

									lowerCamelCase__		:   int  				=					self.get_dummy_inputs(__UpperCAmelCase  )
									lowerCamelCase__		:   int  				=					'''french fries'''
									lowerCamelCase__		:   Any  				=					sd_pipe(**__UpperCAmelCase      ,negative_prompt=__UpperCAmelCase  )
									lowerCamelCase__		:   Optional[Any]  				=					output.images
									lowerCamelCase__		:   str  				=					image[0, -3:, -3:, -1]

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

									lowerCamelCase__		:   Tuple  				=					np.array([0.4_771, 0.5_420, 0.4_683, 0.4_918, 0.5_636, 0.4_725, 0.5_230, 0.4_923, 0.5_015]  )

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





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







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




									"""simple docstring"""



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







				def   lowercase_		(    self    :Optional[int]      ,__UpperCAmelCase    :List[str]=0  )      ->					str:




									"""simple docstring"""



									lowerCamelCase__		:   str  				=					torch.Generator(device=__UpperCAmelCase  ).manual_seed(__UpperCAmelCase  )
									lowerCamelCase__		:   List[Any]  				=					{
									    '''prompt''': '''a photograph of an astronaut riding a horse''',
									    '''generator''': generator,
									    '''num_inference_steps''': 10,
									    '''guidance_scale''': 7.5,
									    '''output_type''': '''numpy''',
									    '''parallel''': 3,
									    '''debug''': True,
									}
									return inputs







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




									"""simple docstring"""



									lowerCamelCase__		:   Optional[Any]  				=					'''stabilityai/stable-diffusion-2-base'''
									lowerCamelCase__		:   Tuple  				=					DDIMParallelScheduler.from_pretrained(__UpperCAmelCase      ,subfolder='''scheduler'''  )
									lowerCamelCase__		:   Optional[Any]  				=					StableDiffusionParadigmsPipeline.from_pretrained(__UpperCAmelCase      ,scheduler=__UpperCAmelCase      ,safety_checker=__UpperCAmelCase  )
									pipe.to(__UpperCAmelCase  )
									pipe.set_progress_bar_config(disable=__UpperCAmelCase  )
									pipe.enable_attention_slicing()

									lowerCamelCase__		:   Optional[Any]  				=					self.get_inputs()
									lowerCamelCase__		:   Tuple  				=					pipe(**__UpperCAmelCase  ).images
									lowerCamelCase__		:   List[Any]  				=					image[0, -3:, -3:, -1].flatten()

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

									lowerCamelCase__		:   List[str]  				=					np.array([0.9_622, 0.9_602, 0.9_748, 0.9_591, 0.9_630, 0.9_691, 0.9_661, 0.9_631, 0.9_741]  )

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






