



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    __lowerCAmelCase					(	lowerCamelCase__			,		lowerCamelCase__			,		unittest.TestCase			):
				__lowerCamelCase					=							StableDiffusionParadigmsPipeline
				__lowerCamelCase					=							TEXT_TO_IMAGE_PARAMS
				__lowerCamelCase					=							TEXT_TO_IMAGE_BATCH_PARAMS
				__lowerCamelCase					=							TEXT_TO_IMAGE_IMAGE_PARAMS
				__lowerCamelCase					=							TEXT_TO_IMAGE_IMAGE_PARAMS




				def 				snake_case  (				self						):



									"""simple docstring"""





									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=_snake_case		,							)
									_lowerCAmelCase         =       DDIMParallelScheduler(
									    beta_start=0.0_0085		,							beta_end=0.012		,							beta_schedule="""scaled_linear"""		,							clip_sample=_snake_case		,							set_alpha_to_one=_snake_case		,							)
									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=128		,							)
									torch.manual_seed(0						)
									_lowerCAmelCase         =       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=1000		,							hidden_act="""gelu"""		,							projection_dim=512		,							)
									_lowerCAmelCase         =       CLIPTextModel(_snake_case						)
									_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		,							_snake_case		,							_snake_case=0						):



									"""simple docstring"""





									if str(_snake_case						).startswith("""mps"""						):
														_lowerCAmelCase         =       torch.manual_seed(_snake_case						)
									else:
														_lowerCAmelCase         =       torch.Generator(device=_snake_case						).manual_seed(_snake_case						)
									_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						):



									"""simple docstring"""





									_lowerCAmelCase         =       """cpu"""  # ensure determinism for the device-dependent torch.Generator
									_lowerCAmelCase         =       self.get_dummy_components()
									_lowerCAmelCase         =       StableDiffusionParadigmsPipeline(**_snake_case						)
									_lowerCAmelCase         =       sd_pipe.to(_snake_case						)
									sd_pipe.set_progress_bar_config(disable=_snake_case						)

									_lowerCAmelCase         =       self.get_dummy_inputs(_snake_case						)
									_lowerCAmelCase         =       sd_pipe(**_snake_case						).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						):



									"""simple docstring"""





									_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(**_snake_case						)
									_lowerCAmelCase         =       sd_pipe.to(_snake_case						)
									sd_pipe.set_progress_bar_config(disable=_snake_case						)

									_lowerCAmelCase         =       self.get_dummy_inputs(_snake_case						)
									_lowerCAmelCase         =       sd_pipe(**_snake_case						).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						):



									"""simple docstring"""





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




				def 				snake_case  (				self						):



									"""simple docstring"""





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




				def 				snake_case  (				self						):



									"""simple docstring"""





									_lowerCAmelCase         =       """cpu"""  # ensure determinism for the device-dependent torch.Generator
									_lowerCAmelCase         =       self.get_dummy_components()
									_lowerCAmelCase         =       StableDiffusionParadigmsPipeline(**_snake_case						)
									_lowerCAmelCase         =       sd_pipe.to(_snake_case						)
									sd_pipe.set_progress_bar_config(disable=_snake_case						)

									_lowerCAmelCase         =       self.get_dummy_inputs(_snake_case						)
									_lowerCAmelCase         =       """french fries"""
									_lowerCAmelCase         =       sd_pipe(**_snake_case		,							negative_prompt=_snake_case						)
									_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    __lowerCAmelCase					(	unittest.TestCase			):




				def 				snake_case  (				self						):



									"""simple docstring"""





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




				def 				snake_case  (				self		,							_snake_case=0						):



									"""simple docstring"""





									_lowerCAmelCase         =       torch.Generator(device=_snake_case						).manual_seed(_snake_case						)
									_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						):



									"""simple docstring"""





									_lowerCAmelCase         =       """stabilityai/stable-diffusion-2-base"""
									_lowerCAmelCase         =       DDIMParallelScheduler.from_pretrained(_snake_case		,							subfolder="""scheduler"""						)
									_lowerCAmelCase         =       StableDiffusionParadigmsPipeline.from_pretrained(_snake_case		,							scheduler=_snake_case		,							safety_checker=_snake_case						)
									pipe.to(_snake_case						)
									pipe.set_progress_bar_config(disable=_snake_case						)
									pipe.enable_attention_slicing()

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

									assert image.shape == (1, 512, 512, 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


__lowerCamelCase     :					Optional[int]				  =						{'''configuration_vilt''': ['''VILT_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''ViltConfig''']}

try:
   if not is_vision_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   __lowerCamelCase     :					Dict				  =						['''ViltFeatureExtractor''']
   __lowerCamelCase     :					Dict				  =						['''ViltImageProcessor''']
   __lowerCamelCase     :					str				  =						['''ViltProcessor''']

try:
   if not is_torch_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   __lowerCamelCase     :					Any				  =						[
       '''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

   __lowerCamelCase     :					List[str]				  =						_LazyModule(__name__, globals()['''__file__'''], _import_structure)
