




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    ):
					_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       lowercase_				(							self   ):







						'''simple docstring'''




						torch.manual_seed(0   )
						lowerCamelCase_				:						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=__lowercase					,			)
						lowerCamelCase_				:						Any				  =				DDIMParallelScheduler(
						    beta_start=0.0_00_85					,			beta_end=0.0_12					,			beta_schedule='scaled_linear'					,			clip_sample=__lowercase					,			set_alpha_to_one=__lowercase					,			)
						torch.manual_seed(0   )
						lowerCamelCase_				:						Optional[Any]				  =				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_				:						int				  =				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=1_000					,			hidden_act='gelu'					,			projection_dim=512					,			)
						lowerCamelCase_				:						int				  =				CLIPTextModel(__lowercase   )
						lowerCamelCase_				:						Optional[Any]				  =				CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'   )

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

					def       lowercase_				(							self					,			__lowercase					,			__lowercase=0   ):







						'''simple docstring'''




						if str(__lowercase   ).startswith('mps'   ):
							lowerCamelCase_				:						Any				  =				torch.manual_seed(__lowercase   )
						else:
							lowerCamelCase_				:						str				  =				torch.Generator(device=__lowercase   ).manual_seed(__lowercase   )
						lowerCamelCase_				:						Optional[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       lowercase_				(							self   ):







						'''simple docstring'''




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

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

						lowerCamelCase_				:						List[Any]				  =				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   ):







						'''simple docstring'''




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

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

						lowerCamelCase_				:						List[str]				  =				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   ):







						'''simple docstring'''




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

					def       lowercase_				(							self   ):







						'''simple docstring'''




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





					def       lowercase_				(							self   ):







						'''simple docstring'''




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

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

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

						lowerCamelCase_				:						List[str]				  =				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 							__SCREAMING_SNAKE_CASE				(						unittest.TestCase    ):

					def       lowercase_				(							self   ):







						'''simple docstring'''




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

					def       lowercase_				(							self					,			__lowercase=0   ):







						'''simple docstring'''




						lowerCamelCase_				:						int				  =				torch.Generator(device=__lowercase   ).manual_seed(__lowercase   )
						lowerCamelCase_				:						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       lowercase_				(							self   ):







						'''simple docstring'''




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

						lowerCamelCase_				:						int				  =				self.get_inputs()
						lowerCamelCase_				:						Optional[int]				  =				pipe(**__lowercase   ).images
						lowerCamelCase_				:						int				  =				image[0, -3:, -3:, -1].flatten()

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

						lowerCamelCase_				:						int				  =				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_tf_available, is_torch_available


__A     						=      {'''configuration_hubert''': ['''HUBERT_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''HubertConfig''']}

try:
 if not is_torch_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 __A     						=      [
     '''HUBERT_PRETRAINED_MODEL_ARCHIVE_LIST''',
     '''HubertForCTC''',
     '''HubertForSequenceClassification''',
     '''HubertModel''',
     '''HubertPreTrainedModel''',
 ]


try:
 if not is_tf_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 __A     						=      [
     '''TF_HUBERT_PRETRAINED_MODEL_ARCHIVE_LIST''',
     '''TFHubertForCTC''',
     '''TFHubertModel''',
     '''TFHubertPreTrainedModel''',
 ]

if TYPE_CHECKING:
 from .configuration_hubert import HUBERT_PRETRAINED_CONFIG_ARCHIVE_MAP, HubertConfig

 try:
  if not is_torch_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_hubert import (
      HUBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
      HubertForCTC,
      HubertForSequenceClassification,
      HubertModel,
      HubertPreTrainedModel,
  )

 try:
  if not is_tf_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_tf_hubert import (
      TF_HUBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
      TFHubertForCTC,
      TFHubertModel,
      TFHubertPreTrainedModel,
  )


else:
 import sys

 __A     						=      _LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)

