



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__	(						snake_case__       ,				snake_case__       ,				unittest.TestCase    ):
					__A  :     Optional[Any]     =    StableDiffusionParadigmsPipeline
					__A  :     Tuple     =    TEXT_TO_IMAGE_PARAMS
					__A  :     int     =    TEXT_TO_IMAGE_BATCH_PARAMS
					__A  :     Dict     =    TEXT_TO_IMAGE_IMAGE_PARAMS
					__A  :     List[str]     =    TEXT_TO_IMAGE_IMAGE_PARAMS







					def 			_lowercase      (							self   :  List[Any]):
											torch.manual_seed(0)
											A__						:				Tuple         =			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=__SCREAMING_SNAKE_CASE      ,		)
											A__						:				Optional[int]         =			DDIMParallelScheduler(
											    beta_start=0.0_00_85      ,		beta_end=0.0_12      ,		beta_schedule="scaled_linear"      ,		clip_sample=__SCREAMING_SNAKE_CASE      ,		set_alpha_to_one=__SCREAMING_SNAKE_CASE      ,		)
											torch.manual_seed(0)
											A__						:				str         =			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)
											A__						:				Union[str, 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=1000      ,		hidden_act="gelu"      ,		projection_dim=512      ,		)
											A__						:				str         =			CLIPTextModel(__SCREAMING_SNAKE_CASE)
											A__						:				Optional[int]         =			CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")

											A__						:				str         =			{
											    "unet": unet,
											    "scheduler": scheduler,
											    "vae": vae,
											    "text_encoder": text_encoder,
											    "tokenizer": tokenizer,
											    "safety_checker": None,
											    "feature_extractor": None,
											}
											return components







					def 			_lowercase      (							self   :  List[str]      ,		_A   :  int      ,		_A   :  Any=0):







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

											A__						:				List[Any]         =			self.get_dummy_inputs(__SCREAMING_SNAKE_CASE)
											A__						:				int         =			sd_pipe(**__SCREAMING_SNAKE_CASE).images
											A__						:				Any         =			image[0, -3:, -3:, -1]
											assert image.shape == (1, 64, 64, 3)

											A__						:				int         =			np.array([0.47_73, 0.54_17, 0.47_23, 0.49_25, 0.56_31, 0.47_52, 0.52_40, 0.49_35, 0.50_23])

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







					def 			_lowercase      (							self   :  str):
											A__						:				Optional[Any]         =			"cpu"  # ensure determinism for the device-dependent torch.Generator
											A__						:				Optional[int]         =			self.get_dummy_components()
											torch.manual_seed(0)
											A__						:				Optional[Any]         =			DDPMParallelScheduler()
											torch.manual_seed(0)
											A__						:				List[str]         =			StableDiffusionParadigmsPipeline(**__SCREAMING_SNAKE_CASE)
											A__						:				str         =			sd_pipe.to(__SCREAMING_SNAKE_CASE)
											sd_pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE)

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

											A__						:				Union[str, Any]         =			np.array([0.35_73, 0.44_20, 0.49_60, 0.47_99, 0.37_96, 0.38_79, 0.48_19, 0.43_65, 0.44_68])

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







					def 			_lowercase      (							self   :  Any):
											super().test_inference_batch_consistent(batch_sizes=[1, 2])







					def 			_lowercase      (							self   :  Union[str, Any]):
											super().test_inference_batch_single_identical(batch_size=2      ,		expected_max_diff=3e-3)







					def 			_lowercase      (							self   :  Tuple):
											A__						:				Dict         =			"cpu"  # ensure determinism for the device-dependent torch.Generator
											A__						:				Tuple         =			self.get_dummy_components()
											A__						:				Union[str, Any]         =			StableDiffusionParadigmsPipeline(**__SCREAMING_SNAKE_CASE)
											A__						:				List[str]         =			sd_pipe.to(__SCREAMING_SNAKE_CASE)
											sd_pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE)

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

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

											A__						:				Any         =			np.array([0.47_71, 0.54_20, 0.46_83, 0.49_18, 0.56_36, 0.47_25, 0.52_30, 0.49_23, 0.50_15])

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




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







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







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

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

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

											A__						:				Optional[Any]         =			np.array([0.96_22, 0.96_02, 0.97_48, 0.95_91, 0.96_30, 0.96_91, 0.96_61, 0.96_31, 0.97_41])

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



from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_sentencepiece_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


__UpperCAmelCase      			=  {
    """configuration_switch_transformers""": [
        """SWITCH_TRANSFORMERS_PRETRAINED_CONFIG_ARCHIVE_MAP""",
        """SwitchTransformersConfig""",
        """SwitchTransformersOnnxConfig""",
    ]
}

try:
     if not is_torch_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     __UpperCAmelCase      			=  [
         """SWITCH_TRANSFORMERS_PRETRAINED_MODEL_ARCHIVE_LIST""",
         """SwitchTransformersEncoderModel""",
         """SwitchTransformersForConditionalGeneration""",
         """SwitchTransformersModel""",
         """SwitchTransformersPreTrainedModel""",
         """SwitchTransformersTop1Router""",
         """SwitchTransformersSparseMLP""",
     ]


if TYPE_CHECKING:
     from .configuration_switch_transformers import (
         SWITCH_TRANSFORMERS_PRETRAINED_CONFIG_ARCHIVE_MAP,
         SwitchTransformersConfig,
         SwitchTransformersOnnxConfig,
     )

     try:
          if not is_torch_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_switch_transformers import (
              SWITCH_TRANSFORMERS_PRETRAINED_MODEL_ARCHIVE_LIST,
              SwitchTransformersEncoderModel,
              SwitchTransformersForConditionalGeneration,
              SwitchTransformersModel,
              SwitchTransformersPreTrainedModel,
              SwitchTransformersSparseMLP,
              SwitchTransformersTopaRouter,
          )


else:
     import sys

     __UpperCAmelCase      			=  _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)

