


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






				'''simple docstring'''



				SCREAMING_SNAKE_CASE__    :      int      =		StableDiffusionParadigmsPipeline
				SCREAMING_SNAKE_CASE__    :      str      =		TEXT_TO_IMAGE_PARAMS
				SCREAMING_SNAKE_CASE__    :      Tuple      =		TEXT_TO_IMAGE_BATCH_PARAMS
				SCREAMING_SNAKE_CASE__    :      int      =		TEXT_TO_IMAGE_IMAGE_PARAMS
				SCREAMING_SNAKE_CASE__    :      Any      =		TEXT_TO_IMAGE_IMAGE_PARAMS


				def 		__magic_name__(	self      :Union[str, Any]     )		->      List[str]:
										torch.manual_seed(0     )
										__SCREAMING_SNAKE_CASE						:					List[Any]		      =		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=lowerCAmelCase__      ,      )
										__SCREAMING_SNAKE_CASE						:					List[str]		      =		DDIMParallelScheduler(
										    beta_start=0.0_0085      ,      beta_end=0.012      ,      beta_schedule='''scaled_linear'''      ,      clip_sample=lowerCAmelCase__      ,      set_alpha_to_one=lowerCAmelCase__      ,      )
										torch.manual_seed(0     )
										__SCREAMING_SNAKE_CASE						:					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=128      ,      )
										torch.manual_seed(0     )
										__SCREAMING_SNAKE_CASE						:					Optional[int]		      =		CLIPTextConfig(
										    bos_token_id=0      ,      eos_token_id=2      ,      hidden_size=32      ,      intermediate_size=37      ,      layer_norm_eps=1E-0_5      ,      num_attention_heads=4      ,      num_hidden_layers=5      ,      pad_token_id=1      ,      vocab_size=1_000      ,      hidden_act='''gelu'''      ,      projection_dim=512      ,      )
										__SCREAMING_SNAKE_CASE						:					int		      =		CLIPTextModel(lowerCAmelCase__     )
										__SCREAMING_SNAKE_CASE						:					str		      =		CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''     )

										__SCREAMING_SNAKE_CASE						:					int		      =		{
										    '''unet''': unet,
										    '''scheduler''': scheduler,
										    '''vae''': vae,
										    '''text_encoder''': text_encoder,
										    '''tokenizer''': tokenizer,
										    '''safety_checker''': None,
										    '''feature_extractor''': None,
										}
										return components


				def 		__magic_name__(	self      :Optional[int]      ,      lowerCAmelCase__      :Tuple      ,      lowerCAmelCase__      :str=0     )		->      Union[str, Any]:
										if str(lowerCAmelCase__     ).startswith('''mps'''     ):
																__SCREAMING_SNAKE_CASE						:					Any		      =		torch.manual_seed(lowerCAmelCase__     )
										else:
																__SCREAMING_SNAKE_CASE						:					List[str]		      =		torch.Generator(device=lowerCAmelCase__     ).manual_seed(lowerCAmelCase__     )
										__SCREAMING_SNAKE_CASE						:					Any		      =		{
										    '''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 		__magic_name__(	self      :List[str]     )		->      Tuple:
										__SCREAMING_SNAKE_CASE						:					List[str]		      =		'''cpu'''  # ensure determinism for the device-dependent torch.Generator
										__SCREAMING_SNAKE_CASE						:					Union[str, Any]		      =		self.get_dummy_components()
										__SCREAMING_SNAKE_CASE						:					Any		      =		StableDiffusionParadigmsPipeline(**lowerCAmelCase__     )
										__SCREAMING_SNAKE_CASE						:					Union[str, Any]		      =		sd_pipe.to(lowerCAmelCase__     )
										sd_pipe.set_progress_bar_config(disable=lowerCAmelCase__     )

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

										__SCREAMING_SNAKE_CASE						:					Union[str, Any]		      =		np.array([0.4773, 0.5417, 0.4723, 0.4925, 0.5631, 0.4752, 0.5240, 0.4935, 0.5023]     )

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


				def 		__magic_name__(	self      :Any     )		->      Dict:
										__SCREAMING_SNAKE_CASE						:					int		      =		'''cpu'''  # ensure determinism for the device-dependent torch.Generator
										__SCREAMING_SNAKE_CASE						:					Optional[int]		      =		self.get_dummy_components()
										torch.manual_seed(0     )
										__SCREAMING_SNAKE_CASE						:					Union[str, Any]		      =		DDPMParallelScheduler()
										torch.manual_seed(0     )
										__SCREAMING_SNAKE_CASE						:					List[str]		      =		StableDiffusionParadigmsPipeline(**lowerCAmelCase__     )
										__SCREAMING_SNAKE_CASE						:					Union[str, Any]		      =		sd_pipe.to(lowerCAmelCase__     )
										sd_pipe.set_progress_bar_config(disable=lowerCAmelCase__     )

										__SCREAMING_SNAKE_CASE						:					Dict		      =		self.get_dummy_inputs(lowerCAmelCase__     )
										__SCREAMING_SNAKE_CASE						:					List[str]		      =		sd_pipe(**lowerCAmelCase__     ).images
										__SCREAMING_SNAKE_CASE						:					Dict		      =		image[0, -3:, -3:, -1]
										assert image.shape == (1, 64, 64, 3)

										__SCREAMING_SNAKE_CASE						:					List[str]		      =		np.array([0.3573, 0.4420, 0.4960, 0.4799, 0.3796, 0.3879, 0.4819, 0.4365, 0.4468]     )

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


				def 		__magic_name__(	self      :int     )		->      Any:
										super().test_inference_batch_consistent(batch_sizes=[1, 2]     )


				def 		__magic_name__(	self      :Tuple     )		->      Optional[Any]:
										super().test_inference_batch_single_identical(batch_size=2      ,      expected_max_diff=3E-3     )



				def 		__magic_name__(	self      :int     )		->      Dict:
										__SCREAMING_SNAKE_CASE						:					List[Any]		      =		'''cpu'''  # ensure determinism for the device-dependent torch.Generator
										__SCREAMING_SNAKE_CASE						:					Dict		      =		self.get_dummy_components()
										__SCREAMING_SNAKE_CASE						:					Optional[Any]		      =		StableDiffusionParadigmsPipeline(**lowerCAmelCase__     )
										__SCREAMING_SNAKE_CASE						:					int		      =		sd_pipe.to(lowerCAmelCase__     )
										sd_pipe.set_progress_bar_config(disable=lowerCAmelCase__     )

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

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

										__SCREAMING_SNAKE_CASE						:					Optional[int]		      =		np.array([0.4771, 0.5420, 0.4683, 0.4918, 0.5636, 0.4725, 0.5230, 0.4923, 0.5015]     )

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




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






				'''simple docstring'''



				def 		__magic_name__(	self      :str     )		->      List[str]:
										super().tearDown()
										gc.collect()
										torch.cuda.empty_cache()


				def 		__magic_name__(	self      :str      ,      lowerCAmelCase__      :Dict=0     )		->      Any:
										__SCREAMING_SNAKE_CASE						:					str		      =		torch.Generator(device=lowerCAmelCase__     ).manual_seed(lowerCAmelCase__     )
										__SCREAMING_SNAKE_CASE						:					Optional[int]		      =		{
										    '''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 		__magic_name__(	self      :Optional[int]     )		->      Any:
										__SCREAMING_SNAKE_CASE						:					List[Any]		      =		'''stabilityai/stable-diffusion-2-base'''
										__SCREAMING_SNAKE_CASE						:					Tuple		      =		DDIMParallelScheduler.from_pretrained(lowerCAmelCase__      ,      subfolder='''scheduler'''     )
										__SCREAMING_SNAKE_CASE						:					Tuple		      =		StableDiffusionParadigmsPipeline.from_pretrained(lowerCAmelCase__      ,      scheduler=lowerCAmelCase__      ,      safety_checker=lowerCAmelCase__     )
										pipe.to(lowerCAmelCase__     )
										pipe.set_progress_bar_config(disable=lowerCAmelCase__     )
										pipe.enable_attention_slicing()

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

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

										__SCREAMING_SNAKE_CASE						:					Optional[Any]		      =		np.array([0.9622, 0.9602, 0.9748, 0.9591, 0.9630, 0.9691, 0.9661, 0.9631, 0.9741]     )

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




from ctypes import c_float, sizeof
from enum import Enum
from typing import TYPE_CHECKING, Optional, Union


if TYPE_CHECKING:
						from .. import AutoFeatureExtractor, AutoProcessor, AutoTokenizer  # tests_ignore
class 	_lowercase				(      A__	):






				'''simple docstring'''



				SCREAMING_SNAKE_CASE__    :      int      =		c_float


				@property
				def 		__magic_name__(	self      :int     )		->      int:
										return sizeof(self.value     )




def 	_UpperCamelCase							(				lowercase__   ,	lowercase__   ,	lowercase__ = 0				):
						# < 0 is possible if using a dynamic axis
						if dimension <= 0:
												__SCREAMING_SNAKE_CASE						:					int		      =		fixed_dimension

						dimension -= num_token_to_add
						return dimension


def 	_UpperCamelCase							(				lowercase__   ,	lowercase__				):
						return num_parameters * dtype.size


def 	_UpperCamelCase							(				lowercase__				):
						from .. import AutoFeatureExtractor, AutoProcessor, AutoTokenizer  # tests_ignore

						try:
												return AutoProcessor.from_pretrained(lowercase__				)
						except (ValueError, OSError, KeyError):
												__SCREAMING_SNAKE_CASE      ,      __SCREAMING_SNAKE_CASE						:					Tuple		      =		None, None
												try:
																		__SCREAMING_SNAKE_CASE						:					List[Any]		      =		AutoTokenizer.from_pretrained(lowercase__				)
												except (OSError, KeyError):
																		pass
												try:
																		__SCREAMING_SNAKE_CASE						:					Optional[int]		      =		AutoFeatureExtractor.from_pretrained(lowercase__				)
												except (OSError, KeyError):
																		pass

												if tokenizer is not None and feature_extractor is not None:
																		raise ValueError(
																		    F'''Couldn\'t auto-detect preprocessor for {model_name}. Found both a tokenizer and a feature extractor.'''				)
												elif tokenizer is None and feature_extractor is None:
																		return None
												elif tokenizer is not None:
																		return tokenizer
												else:
																		return feature_extractor
