





'''simple docstring'''






import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    StableDiffusionSAGPipeline,
    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   (				lowercase_   ,     lowercase_   ,     unittest.TestCase):


  """simple docstring"""







  lowerCAmelCase_					=				StableDiffusionSAGPipeline
  lowerCAmelCase_					=				TEXT_TO_IMAGE_PARAMS
  lowerCAmelCase_					=				TEXT_TO_IMAGE_BATCH_PARAMS
  lowerCAmelCase_					=				TEXT_TO_IMAGE_IMAGE_PARAMS
  lowerCAmelCase_					=				TEXT_TO_IMAGE_IMAGE_PARAMS
  lowerCAmelCase_					=				False


  def   UpperCamelCase__      (				self    :					List[Any]  )     ->						Optional[Any]:
    torch.manual_seed(0  )
    _UpperCamelCase       			=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							,	)
    _UpperCamelCase       			=DDIMScheduler(
        beta_start=0.00085							,	beta_end=0.012							,	beta_schedule='''scaled_linear'''							,	clip_sample=UpperCamelCase__							,	set_alpha_to_one=UpperCamelCase__							,	)
    torch.manual_seed(0  )
    _UpperCamelCase       			=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							,	)
    torch.manual_seed(0  )
    _UpperCamelCase       			=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							,	)
    _UpperCamelCase       			=CLIPTextModel(UpperCamelCase__  )
    _UpperCamelCase       			=CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''  )

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


  def   UpperCamelCase__      (				self    :					Optional[int]							,	UpperCamelCase__    :					Dict							,	UpperCamelCase__    :					Optional[int]=0  )     ->						str:
    if str(UpperCamelCase__  ).startswith('''mps'''  ):
      _UpperCamelCase       			=torch.manual_seed(UpperCamelCase__  )
    else:
      _UpperCamelCase       			=torch.Generator(device=UpperCamelCase__  ).manual_seed(UpperCamelCase__  )
    _UpperCamelCase       			={
        '''prompt''': '''.''',
        '''generator''': generator,
        '''num_inference_steps''': 2,
        '''guidance_scale''': 1.0,
        '''sag_scale''': 1.0,
        '''output_type''': '''numpy''',
    }
    return inputs




  def   UpperCamelCase__      (				self    :					int  )     ->						Tuple:
    super().test_inference_batch_single_identical(expected_max_diff=3E-3  )







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


  """simple docstring"""







  def   UpperCamelCase__      (				self    :					Optional[int]  )     ->						Optional[int]:
    # clean up the VRAM after each test
    super().tearDown()
    gc.collect()
    torch.cuda.empty_cache()


  def   UpperCamelCase__      (				self    :					Any  )     ->						Optional[Any]:
    _UpperCamelCase       			=StableDiffusionSAGPipeline.from_pretrained('''CompVis/stable-diffusion-v1-4'''  )
    _UpperCamelCase       			=sag_pipe.to(UpperCamelCase__  )
    sag_pipe.set_progress_bar_config(disable=UpperCamelCase__  )

    _UpperCamelCase       			='''.'''
    _UpperCamelCase       			=torch.manual_seed(0  )
    _UpperCamelCase       			=sag_pipe(
        [prompt]							,	generator=UpperCamelCase__							,	guidance_scale=7.5							,	sag_scale=1.0							,	num_inference_steps=20							,	output_type='''np'''  )

    _UpperCamelCase       			=output.images

    _UpperCamelCase       			=image[0, -3:, -3:, -1]

    assert image.shape == (1, 512, 512, 3)
    _UpperCamelCase       			=np.array([0.1568, 0.1738, 0.1695, 0.1693, 0.1507, 0.1705, 0.1547, 0.1751, 0.1949]  )

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


  def   UpperCamelCase__      (				self    :					Union[str, Any]  )     ->						Tuple:
    _UpperCamelCase       			=StableDiffusionSAGPipeline.from_pretrained('''stabilityai/stable-diffusion-2-1-base'''  )
    _UpperCamelCase       			=sag_pipe.to(UpperCamelCase__  )
    sag_pipe.set_progress_bar_config(disable=UpperCamelCase__  )

    _UpperCamelCase       			='''.'''
    _UpperCamelCase       			=torch.manual_seed(0  )
    _UpperCamelCase       			=sag_pipe(
        [prompt]							,	generator=UpperCamelCase__							,	guidance_scale=7.5							,	sag_scale=1.0							,	num_inference_steps=20							,	output_type='''np'''  )

    _UpperCamelCase       			=output.images

    _UpperCamelCase       			=image[0, -3:, -3:, -1]

    assert image.shape == (1, 512, 512, 3)
    _UpperCamelCase       			=np.array([0.3459, 0.2876, 0.2537, 0.3002, 0.2671, 0.2160, 0.3026, 0.2262, 0.2371]  )

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




  def   UpperCamelCase__      (				self    :					str  )     ->						Tuple:
    _UpperCamelCase       			=StableDiffusionSAGPipeline.from_pretrained('''stabilityai/stable-diffusion-2-1-base'''  )
    _UpperCamelCase       			=sag_pipe.to(UpperCamelCase__  )
    sag_pipe.set_progress_bar_config(disable=UpperCamelCase__  )

    _UpperCamelCase       			='''.'''
    _UpperCamelCase       			=torch.manual_seed(0  )
    _UpperCamelCase       			=sag_pipe(
        [prompt]							,	width=768							,	height=512							,	generator=UpperCamelCase__							,	guidance_scale=7.5							,	sag_scale=1.0							,	num_inference_steps=20							,	output_type='''np'''							,	)

    _UpperCamelCase       			=output.images

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













'''simple docstring'''






from typing import TYPE_CHECKING

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


__lowerCamelCase						: str		=   {
    'configuration_speech_to_text': ['SPEECH_TO_TEXT_PRETRAINED_CONFIG_ARCHIVE_MAP', 'Speech2TextConfig'],
    'processing_speech_to_text': ['Speech2TextProcessor'],
}

try:
      if not is_sentencepiece_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      __lowerCamelCase						: List[str]		=   ['Speech2TextTokenizer']

try:
      if not is_speech_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      __lowerCamelCase						: str		=   ['Speech2TextFeatureExtractor']

try:
      if not is_tf_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      __lowerCamelCase						: Optional[Any]		=   [
          'TF_SPEECH_TO_TEXT_PRETRAINED_MODEL_ARCHIVE_LIST',
          'TFSpeech2TextForConditionalGeneration',
          'TFSpeech2TextModel',
          'TFSpeech2TextPreTrainedModel',
      ]

try:
      if not is_torch_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      __lowerCamelCase						: Union[str, Any]		=   [
          'SPEECH_TO_TEXT_PRETRAINED_MODEL_ARCHIVE_LIST',
          'Speech2TextForConditionalGeneration',
          'Speech2TextModel',
          'Speech2TextPreTrainedModel',
      ]


if TYPE_CHECKING:
      from .configuration_speech_to_text import SPEECH_TO_TEXT_PRETRAINED_CONFIG_ARCHIVE_MAP, SpeechaTextConfig
      from .processing_speech_to_text import SpeechaTextProcessor

      try:
            if not is_sentencepiece_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .tokenization_speech_to_text import SpeechaTextTokenizer

      try:
            if not is_speech_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .feature_extraction_speech_to_text import SpeechaTextFeatureExtractor

      try:
            if not is_tf_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .modeling_tf_speech_to_text import (
                TF_SPEECH_TO_TEXT_PRETRAINED_MODEL_ARCHIVE_LIST,
                TFSpeechaTextForConditionalGeneration,
                TFSpeechaTextModel,
                TFSpeechaTextPreTrainedModel,
            )

      try:
            if not is_torch_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .modeling_speech_to_text import (
                SPEECH_TO_TEXT_PRETRAINED_MODEL_ARCHIVE_LIST,
                SpeechaTextForConditionalGeneration,
                SpeechaTextModel,
                SpeechaTextPreTrainedModel,
            )

else:
      import sys

      __lowerCamelCase						: Optional[int]		=   _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)






