


import gc
import unittest

import numpy as np
import torch

from diffusers import StableDiffusionKDiffusionPipeline
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()




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






      def        lowerCamelCase       (					self						:     Dict )					->						List[Any]:
             """simple docstring"""




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






      def        lowerCamelCase       (					self						:     Tuple )					->						int:
             """simple docstring"""




             _lowerCamelCase       :		Optional[Any]		=StableDiffusionKDiffusionPipeline.from_pretrained('CompVis/stable-diffusion-v1-4' )
             _lowerCamelCase       :		Optional[int]		=sd_pipe.to(lowercase_ )
             sd_pipe.set_progress_bar_config(disable=lowercase_ )

             sd_pipe.set_scheduler('sample_euler' )

             _lowerCamelCase       :		List[Any]		='A painting of a squirrel eating a burger'
             _lowerCamelCase       :		str		=torch.manual_seed(0 )
             _lowerCamelCase       :		List[Any]		=sd_pipe([prompt]			,   generator=lowercase_			,   guidance_scale=9.0			,   num_inference_steps=20			,   output_type='np' )

             _lowerCamelCase       :		int		=output.images

             _lowerCamelCase       :		Dict		=image[0, -3:, -3:, -1]

             assert image.shape == (1, 512, 512, 3)
             _lowerCamelCase       :		Optional[Any]		=np.array([0.0447, 0.0492, 0.0468, 0.0408, 0.0383, 0.0408, 0.0354, 0.0380, 0.0339] )

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






      def        lowerCamelCase       (					self						:     List[str] )					->						List[Any]:
             """simple docstring"""




             _lowerCamelCase       :		Dict		=StableDiffusionKDiffusionPipeline.from_pretrained('stabilityai/stable-diffusion-2-1-base' )
             _lowerCamelCase       :		Dict		=sd_pipe.to(lowercase_ )
             sd_pipe.set_progress_bar_config(disable=lowercase_ )

             sd_pipe.set_scheduler('sample_euler' )

             _lowerCamelCase       :		int		='A painting of a squirrel eating a burger'
             _lowerCamelCase       :		List[str]		=torch.manual_seed(0 )
             _lowerCamelCase       :		str		=sd_pipe([prompt]			,   generator=lowercase_			,   guidance_scale=9.0			,   num_inference_steps=20			,   output_type='np' )

             _lowerCamelCase       :		Optional[Any]		=output.images

             _lowerCamelCase       :		int		=image[0, -3:, -3:, -1]

             assert image.shape == (1, 512, 512, 3)
             _lowerCamelCase       :		str		=np.array([0.1237, 0.1320, 0.1438, 0.1359, 0.1390, 0.1132, 0.1277, 0.1175, 0.1112] )

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






      def        lowerCamelCase       (					self						:     Dict )					->						List[str]:
             """simple docstring"""




             _lowerCamelCase       :		Optional[int]		=StableDiffusionKDiffusionPipeline.from_pretrained('stabilityai/stable-diffusion-2-1-base' )
             _lowerCamelCase       :		Any		=sd_pipe.to(lowercase_ )
             sd_pipe.set_progress_bar_config(disable=lowercase_ )

             sd_pipe.set_scheduler('sample_dpmpp_2m' )

             _lowerCamelCase       :		Optional[int]		='A painting of a squirrel eating a burger'
             _lowerCamelCase       :		Optional[int]		=torch.manual_seed(0 )
             _lowerCamelCase       :		List[Any]		=sd_pipe(
                 [prompt]			,   generator=lowercase_			,   guidance_scale=7.5			,   num_inference_steps=15			,   output_type='np'			,   use_karras_sigmas=lowercase_			,   )

             _lowerCamelCase       :		Dict		=output.images

             _lowerCamelCase       :		str		=image[0, -3:, -3:, -1]

             assert image.shape == (1, 512, 512, 3)
             _lowerCamelCase       :		Optional[int]		=np.array(
                 [0.11381689, 0.12112921, 0.1389457, 0.12549606, 0.1244964, 0.10831517, 0.11562866, 0.10867816, 0.10499048] )

             assert np.abs(image_slice.flatten() - expected_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,
)


if is_sentencepiece_available():
 from ..ta.tokenization_ta import TaTokenizer
else:
 from ...utils.dummy_sentencepiece_objects import TaTokenizer

lowerCamelCase											=       TaTokenizer

if is_tokenizers_available():
 from ..ta.tokenization_ta_fast import TaTokenizerFast
else:
 from ...utils.dummy_tokenizers_objects import TaTokenizerFast

lowerCamelCase											=       TaTokenizerFast

lowerCamelCase											=       {'configuration_mt5': ['MT5Config', 'MT5OnnxConfig']}

try:
 if not is_torch_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 lowerCamelCase											=       [
     'MT5EncoderModel',
     'MT5ForConditionalGeneration',
     'MT5ForQuestionAnswering',
     'MT5Model',
     'MT5PreTrainedModel',
     'MT5Stack',
 ]

try:
 if not is_tf_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 lowerCamelCase											=       ['TFMT5EncoderModel', 'TFMT5ForConditionalGeneration', 'TFMT5Model']

try:
 if not is_flax_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 lowerCamelCase											=       ['FlaxMT5EncoderModel', 'FlaxMT5ForConditionalGeneration', 'FlaxMT5Model']


if TYPE_CHECKING:
 from .configuration_mta import MTaConfig, MTaOnnxConfig

 try:
  if not is_torch_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_mta import (
      MTaEncoderModel,
      MTaForConditionalGeneration,
      MTaForQuestionAnswering,
      MTaModel,
      MTaPreTrainedModel,
      MTaStack,
  )

 try:
  if not is_tf_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_tf_mta import TFMTaEncoderModel, TFMTaForConditionalGeneration, TFMTaModel

 try:
  if not is_flax_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_flax_mta import FlaxMTaEncoderModel, FlaxMTaForConditionalGeneration, FlaxMTaModel

else:
 import sys

 lowerCamelCase											=       _LazyModule(
     __name__,
     globals()['__file__'],
     _import_structure,
     extra_objects={'MT5Tokenizer': MTaTokenizer, 'MT5TokenizerFast': MTaTokenizerFast},
     module_spec=__spec__,
 )


