
"""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				(		__a  ,      __a  ,      unittest.TestCase				):
      __SCREAMING_SNAKE_CASE :				List[str]				   =						StableDiffusionSAGPipeline
      __SCREAMING_SNAKE_CASE :				List[str]				   =						TEXT_TO_IMAGE_PARAMS
      __SCREAMING_SNAKE_CASE :				Union[str, Any]				   =						TEXT_TO_IMAGE_BATCH_PARAMS
      __SCREAMING_SNAKE_CASE :				Tuple				   =						TEXT_TO_IMAGE_IMAGE_PARAMS
      __SCREAMING_SNAKE_CASE :				Dict				   =						TEXT_TO_IMAGE_IMAGE_PARAMS
      __SCREAMING_SNAKE_CASE :				str				   =						False




      def      a_   (				self				)    ->       List[str]:
          torch.manual_seed(0				)
          UpperCAmelCase          =       UNetaDConditionModel(
              block_out_channels=(3_2, 6_4) ,							layers_per_block=2 ,							sample_size=3_2 ,							in_channels=4 ,							out_channels=4 ,							down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D') ,							up_block_types=('CrossAttnUpBlock2D', 'UpBlock2D') ,							cross_attention_dim=3_2 ,							)
          UpperCAmelCase          =       DDIMScheduler(
              beta_start=0.0_0_0_8_5 ,							beta_end=0.0_1_2 ,							beta_schedule='scaled_linear' ,							clip_sample=_lowerCamelCase ,							set_alpha_to_one=_lowerCamelCase ,							)
          torch.manual_seed(0				)
          UpperCAmelCase          =       AutoencoderKL(
              block_out_channels=[3_2, 6_4] ,							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=3_2 ,							intermediate_size=3_7 ,							layer_norm_eps=1E-05 ,							num_attention_heads=4 ,							num_hidden_layers=5 ,							pad_token_id=1 ,							vocab_size=1_0_0_0 ,							)
          UpperCAmelCase          =       CLIPTextModel(_lowerCamelCase				)
          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      a_   (				self ,							lowercase_ ,							lowercase_=0				)    ->       int:
          if str(_lowerCamelCase				).startswith('mps'				):
              UpperCAmelCase          =       torch.manual_seed(_lowerCamelCase				)
          else:
              UpperCAmelCase          =       torch.Generator(device=_lowerCamelCase				).manual_seed(_lowerCamelCase				)
          UpperCAmelCase          =       {
              'prompt': '.',
              'generator': generator,
              'num_inference_steps': 2,
              'guidance_scale': 1.0,
              'sag_scale': 1.0,
              'output_type': 'numpy',
          }
          return inputs




      def      a_   (				self				)    ->       List[Any]:
          super().test_inference_batch_single_identical(expected_max_diff=3E-3				)

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




      def      a_   (				self				)    ->       Dict:
          super().tearDown()
          gc.collect()
          torch.cuda.empty_cache()




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

          UpperCAmelCase          =       '.'
          UpperCAmelCase          =       torch.manual_seed(0				)
          UpperCAmelCase          =       sag_pipe(
              [prompt] ,							generator=_lowerCamelCase ,							guidance_scale=7.5 ,							sag_scale=1.0 ,							num_inference_steps=2_0 ,							output_type='np'				)

          UpperCAmelCase          =       output.images

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

          assert image.shape == (1, 5_1_2, 5_1_2, 3)
          UpperCAmelCase          =       np.array([0.1_5_6_8, 0.1_7_3_8, 0.1_6_9_5, 0.1_6_9_3, 0.1_5_0_7, 0.1_7_0_5, 0.1_5_4_7, 0.1_7_5_1, 0.1_9_4_9]				)

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




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

          UpperCAmelCase          =       '.'
          UpperCAmelCase          =       torch.manual_seed(0				)
          UpperCAmelCase          =       sag_pipe(
              [prompt] ,							generator=_lowerCamelCase ,							guidance_scale=7.5 ,							sag_scale=1.0 ,							num_inference_steps=2_0 ,							output_type='np'				)

          UpperCAmelCase          =       output.images

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

          assert image.shape == (1, 5_1_2, 5_1_2, 3)
          UpperCAmelCase          =       np.array([0.3_4_5_9, 0.2_8_7_6, 0.2_5_3_7, 0.3_0_0_2, 0.2_6_7_1, 0.2_1_6_0, 0.3_0_2_6, 0.2_2_6_2, 0.2_3_7_1]				)

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




      def      a_   (				self				)    ->       List[Any]:
          UpperCAmelCase          =       StableDiffusionSAGPipeline.from_pretrained('stabilityai/stable-diffusion-2-1-base'				)
          UpperCAmelCase          =       sag_pipe.to(_lowerCamelCase				)
          sag_pipe.set_progress_bar_config(disable=_lowerCamelCase				)

          UpperCAmelCase          =       '.'
          UpperCAmelCase          =       torch.manual_seed(0				)
          UpperCAmelCase          =       sag_pipe(
              [prompt] ,							width=7_6_8 ,							height=5_1_2 ,							generator=_lowerCamelCase ,							guidance_scale=7.5 ,							sag_scale=1.0 ,							num_inference_steps=2_0 ,							output_type='np' ,							)

          UpperCAmelCase          =       output.images

          assert image.shape == (1, 5_1_2, 7_6_8, 3)





"""simple docstring"""






import warnings
from contextlib import contextmanager

from ...processing_utils import ProcessorMixin
from .feature_extraction_wavaveca import WavaVecaFeatureExtractor
from .tokenization_wavaveca import WavaVecaCTCTokenizer

class 					__a     ( __a       ):



							'''simple docstring'''







							_lowerCamelCase					:			List[Any]          =   """Wav2Vec2FeatureExtractor"""
							_lowerCamelCase					:			Any          =   """AutoTokenizer"""
							def __init__(     self  ,    _lowerCamelCase  ,    _lowerCamelCase       )			->							List[str]:




										'''simple docstring'''
										super().__init__(_lowerCamelCase  ,    _lowerCamelCase       )
										__lowercase       =       self.feature_extractor
										__lowercase       =       False
							@classmethod
							def   SCREAMING_SNAKE_CASE						(     cls  ,    _lowerCamelCase  ,    **_lowerCamelCase       )			->							Any:




										'''simple docstring'''
										try:
													return super().from_pretrained(_lowerCamelCase  ,    **_lowerCamelCase       )
										except OSError:
													warnings.warn(
													    f'''Loading a tokenizer inside {cls.__name__} from a config that does not'''
													    " include a `tokenizer_class` attribute is deprecated and will be "
													    "removed in v5. Please add `'tokenizer_class': 'Wav2Vec2CTCTokenizer'`"
													    " attribute to either your `config.json` or `tokenizer_config.json` "
													    "file to suppress this warning: "  ,    _lowerCamelCase  ,    )

													__lowercase       =       WavaVecaFeatureExtractor.from_pretrained(_lowerCamelCase  ,    **_lowerCamelCase       )
													__lowercase       =       WavaVecaCTCTokenizer.from_pretrained(_lowerCamelCase  ,    **_lowerCamelCase       )

													return cls(feature_extractor=_lowerCamelCase  ,    tokenizer=_lowerCamelCase       )
							def __call__(     self  ,    *_lowerCamelCase  ,    **_lowerCamelCase       )			->							List[Any]:




										'''simple docstring'''
										# For backward compatibility
										if self._in_target_context_manager:
													return self.current_processor(*_lowerCamelCase  ,    **_lowerCamelCase       )

										if "raw_speech" in kwargs:
													warnings.warn("Using `raw_speech` as a keyword argument is deprecated. Use `audio` instead."       )
													__lowercase       =       kwargs.pop("raw_speech"       )
										else:
													__lowercase       =       kwargs.pop("audio"  ,    _lowerCamelCase       )
										__lowercase       =       kwargs.pop("sampling_rate"  ,    _lowerCamelCase       )
										__lowercase       =       kwargs.pop("text"  ,    _lowerCamelCase       )
										if len(_lowerCamelCase       ) > 0:
													__lowercase       =       args[0]
													__lowercase       =       args[1:]

										if audio is None and text is None:
													raise ValueError("You need to specify either an `audio` or `text` input to process."       )

										if audio is not None:
													__lowercase       =       self.feature_extractor(_lowerCamelCase  ,    *_lowerCamelCase  ,    sampling_rate=_lowerCamelCase  ,    **_lowerCamelCase       )
										if text is not None:
													__lowercase       =       self.tokenizer(_lowerCamelCase  ,    **_lowerCamelCase       )

										if text is None:
													return inputs
										elif audio is None:
													return encodings
										else:
													__lowercase       =       encodings["input_ids"]
													return inputs
							def   SCREAMING_SNAKE_CASE						(     self  ,    *_lowerCamelCase  ,    **_lowerCamelCase       )			->							List[str]:




										'''simple docstring'''
										# For backward compatibility
										if self._in_target_context_manager:
													return self.current_processor.pad(*_lowerCamelCase  ,    **_lowerCamelCase       )

										__lowercase       =       kwargs.pop("input_features"  ,    _lowerCamelCase       )
										__lowercase       =       kwargs.pop("labels"  ,    _lowerCamelCase       )
										if len(_lowerCamelCase       ) > 0:
													__lowercase       =       args[0]
													__lowercase       =       args[1:]

										if input_features is not None:
													__lowercase       =       self.feature_extractor.pad(_lowerCamelCase  ,    *_lowerCamelCase  ,    **_lowerCamelCase       )
										if labels is not None:
													__lowercase       =       self.tokenizer.pad(_lowerCamelCase  ,    **_lowerCamelCase       )

										if labels is None:
													return input_features
										elif input_features is None:
													return labels
										else:
													__lowercase       =       labels["input_ids"]
													return input_features
							def   SCREAMING_SNAKE_CASE						(     self  ,    *_lowerCamelCase  ,    **_lowerCamelCase       )			->							Optional[int]:




										'''simple docstring'''
										return self.tokenizer.batch_decode(*_lowerCamelCase  ,    **_lowerCamelCase       )
							def   SCREAMING_SNAKE_CASE						(     self  ,    *_lowerCamelCase  ,    **_lowerCamelCase       )			->							Dict:




										'''simple docstring'''
										return self.tokenizer.decode(*_lowerCamelCase  ,    **_lowerCamelCase       )





							@contextmanager
							def   SCREAMING_SNAKE_CASE						(     self       )			->							str:




										'''simple docstring'''
										warnings.warn(
										    "`as_target_processor` is deprecated and will be removed in v5 of Transformers. You can process your "
										    "labels by using the argument `text` of the regular `__call__` method (either in the same call as "
										    "your audio inputs, or in a separate call."       )
										__lowercase       =       True
										__lowercase       =       self.tokenizer
										yield
										__lowercase       =       self.feature_extractor
										__lowercase       =       False





