




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 				__UpperCamelCase      (			__UpperCAmelCase  ,       __UpperCAmelCase  ,       unittest.TestCase	):





				'''simple docstring'''



				__a				:			List[Any]					    =StableDiffusionParadigmsPipeline
				__a				:			Optional[Any]					    =TEXT_TO_IMAGE_PARAMS
				__a				:			Any					    =TEXT_TO_IMAGE_BATCH_PARAMS
				__a				:			int					    =TEXT_TO_IMAGE_IMAGE_PARAMS
				__a				:			int					    =TEXT_TO_IMAGE_IMAGE_PARAMS



				def    __snake_case		(			self    ):
											torch.manual_seed(0    )
											lowerCAmelCase =		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 =		DDIMParallelScheduler(
											    beta_start=0.00085							,			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=[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 =		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=10_00							,			hidden_act='''gelu'''							,			projection_dim=5_12							,			)
											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    __snake_case		(			self							,			UpperCAmelCase_							,			UpperCAmelCase_=0    ):
											if str(UpperCAmelCase_    ).startswith('''mps'''    ):
																		lowerCAmelCase =		torch.manual_seed(UpperCAmelCase_    )
											else:
																		lowerCAmelCase =		torch.Generator(device=UpperCAmelCase_    ).manual_seed(UpperCAmelCase_    )
											lowerCAmelCase =		{
											    '''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    __snake_case		(			self    ):
											lowerCAmelCase =		'''cpu'''  # ensure determinism for the device-dependent torch.Generator
											lowerCAmelCase =		self.get_dummy_components()
											lowerCAmelCase =		StableDiffusionParadigmsPipeline(**UpperCAmelCase_    )
											lowerCAmelCase =		sd_pipe.to(UpperCAmelCase_    )
											sd_pipe.set_progress_bar_config(disable=UpperCAmelCase_    )

											lowerCAmelCase =		self.get_dummy_inputs(UpperCAmelCase_    )
											lowerCAmelCase =		sd_pipe(**UpperCAmelCase_    ).images
											lowerCAmelCase =		image[0, -3:, -3:, -1]
											assert image.shape == (1, 64, 64, 3)

											lowerCAmelCase =		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    __snake_case		(			self    ):
											lowerCAmelCase =		'''cpu'''  # ensure determinism for the device-dependent torch.Generator
											lowerCAmelCase =		self.get_dummy_components()
											torch.manual_seed(0    )
											lowerCAmelCase =		DDPMParallelScheduler()
											torch.manual_seed(0    )
											lowerCAmelCase =		StableDiffusionParadigmsPipeline(**UpperCAmelCase_    )
											lowerCAmelCase =		sd_pipe.to(UpperCAmelCase_    )
											sd_pipe.set_progress_bar_config(disable=UpperCAmelCase_    )

											lowerCAmelCase =		self.get_dummy_inputs(UpperCAmelCase_    )
											lowerCAmelCase =		sd_pipe(**UpperCAmelCase_    ).images
											lowerCAmelCase =		image[0, -3:, -3:, -1]
											assert image.shape == (1, 64, 64, 3)

											lowerCAmelCase =		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    __snake_case		(			self    ):
											super().test_inference_batch_consistent(batch_sizes=[1, 2]    )



				def    __snake_case		(			self    ):
											super().test_inference_batch_single_identical(batch_size=2							,			expected_max_diff=3E-3    )






				def    __snake_case		(			self    ):
											lowerCAmelCase =		'''cpu'''  # ensure determinism for the device-dependent torch.Generator
											lowerCAmelCase =		self.get_dummy_components()
											lowerCAmelCase =		StableDiffusionParadigmsPipeline(**UpperCAmelCase_    )
											lowerCAmelCase =		sd_pipe.to(UpperCAmelCase_    )
											sd_pipe.set_progress_bar_config(disable=UpperCAmelCase_    )

											lowerCAmelCase =		self.get_dummy_inputs(UpperCAmelCase_    )
											lowerCAmelCase =		'''french fries'''
											lowerCAmelCase =		sd_pipe(**UpperCAmelCase_							,			negative_prompt=UpperCAmelCase_    )
											lowerCAmelCase =		output.images
											lowerCAmelCase =		image[0, -3:, -3:, -1]

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

											lowerCAmelCase =		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 				__UpperCamelCase      (			unittest.TestCase	):





				'''simple docstring'''



				def    __snake_case		(			self    ):
											super().tearDown()
											gc.collect()
											torch.cuda.empty_cache()



				def    __snake_case		(			self							,			UpperCAmelCase_=0    ):
											lowerCAmelCase =		torch.Generator(device=UpperCAmelCase_    ).manual_seed(UpperCAmelCase_    )
											lowerCAmelCase =		{
											    '''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    __snake_case		(			self    ):
											lowerCAmelCase =		'''stabilityai/stable-diffusion-2-base'''
											lowerCAmelCase =		DDIMParallelScheduler.from_pretrained(UpperCAmelCase_							,			subfolder='''scheduler'''    )
											lowerCAmelCase =		StableDiffusionParadigmsPipeline.from_pretrained(UpperCAmelCase_							,			scheduler=UpperCAmelCase_							,			safety_checker=UpperCAmelCase_    )
											pipe.to(UpperCAmelCase_    )
											pipe.set_progress_bar_config(disable=UpperCAmelCase_    )
											pipe.enable_attention_slicing()

											lowerCAmelCase =		self.get_inputs()
											lowerCAmelCase =		pipe(**UpperCAmelCase_    ).images
											lowerCAmelCase =		image[0, -3:, -3:, -1].flatten()

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

											lowerCAmelCase =		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 typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available, is_vision_available


UpperCAmelCase_    ={"""configuration_vilt""": ["""VILT_PRETRAINED_CONFIG_ARCHIVE_MAP""", """ViltConfig"""]}

try:
		if not is_vision_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		UpperCAmelCase_    =["""ViltFeatureExtractor"""]
		UpperCAmelCase_    =["""ViltImageProcessor"""]
		UpperCAmelCase_    =["""ViltProcessor"""]

try:
		if not is_torch_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		UpperCAmelCase_    =[
		    """VILT_PRETRAINED_MODEL_ARCHIVE_LIST""",
		    """ViltForImageAndTextRetrieval""",
		    """ViltForImagesAndTextClassification""",
		    """ViltForTokenClassification""",
		    """ViltForMaskedLM""",
		    """ViltForQuestionAnswering""",
		    """ViltLayer""",
		    """ViltModel""",
		    """ViltPreTrainedModel""",
		]


if TYPE_CHECKING:
		from .configuration_vilt import VILT_PRETRAINED_CONFIG_ARCHIVE_MAP, ViltConfig

		try:
				if not is_vision_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .feature_extraction_vilt import ViltFeatureExtractor
				from .image_processing_vilt import ViltImageProcessor
				from .processing_vilt import ViltProcessor

		try:
				if not is_torch_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .modeling_vilt import (
				    VILT_PRETRAINED_MODEL_ARCHIVE_LIST,
				    ViltForImageAndTextRetrieval,
				    ViltForImagesAndTextClassification,
				    ViltForMaskedLM,
				    ViltForQuestionAnswering,
				    ViltForTokenClassification,
				    ViltLayer,
				    ViltModel,
				    ViltPreTrainedModel,
				)


else:
		import sys

		UpperCAmelCase_    =_LazyModule(__name__, globals()["""__file__"""], _import_structure)



