'''simple docstring'''





from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tf_available, is_torch_available


_lowerCAmelCase  :List[Any]  				=    {
    """configuration_transfo_xl""": ["""TRANSFO_XL_PRETRAINED_CONFIG_ARCHIVE_MAP""", """TransfoXLConfig"""],
    """tokenization_transfo_xl""": ["""TransfoXLCorpus""", """TransfoXLTokenizer"""],
}

try:
	if not is_torch_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	_lowerCAmelCase  :Union[str, Any]  				=    [
	    """TRANSFO_XL_PRETRAINED_MODEL_ARCHIVE_LIST""",
	    """AdaptiveEmbedding""",
	    """TransfoXLForSequenceClassification""",
	    """TransfoXLLMHeadModel""",
	    """TransfoXLModel""",
	    """TransfoXLPreTrainedModel""",
	    """load_tf_weights_in_transfo_xl""",
	]

try:
	if not is_tf_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	_lowerCAmelCase  :Dict  				=    [
	    """TF_TRANSFO_XL_PRETRAINED_MODEL_ARCHIVE_LIST""",
	    """TFAdaptiveEmbedding""",
	    """TFTransfoXLForSequenceClassification""",
	    """TFTransfoXLLMHeadModel""",
	    """TFTransfoXLMainLayer""",
	    """TFTransfoXLModel""",
	    """TFTransfoXLPreTrainedModel""",
	]


if TYPE_CHECKING:
	from .configuration_transfo_xl import TRANSFO_XL_PRETRAINED_CONFIG_ARCHIVE_MAP, TransfoXLConfig
	from .tokenization_transfo_xl import TransfoXLCorpus, TransfoXLTokenizer

	try:
		if not is_torch_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .modeling_transfo_xl import (
		    TRANSFO_XL_PRETRAINED_MODEL_ARCHIVE_LIST,
		    AdaptiveEmbedding,
		    TransfoXLForSequenceClassification,
		    TransfoXLLMHeadModel,
		    TransfoXLModel,
		    TransfoXLPreTrainedModel,
		    load_tf_weights_in_transfo_xl,
		)

	try:
		if not is_tf_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .modeling_tf_transfo_xl import (
		    TF_TRANSFO_XL_PRETRAINED_MODEL_ARCHIVE_LIST,
		    TFAdaptiveEmbedding,
		    TFTransfoXLForSequenceClassification,
		    TFTransfoXLLMHeadModel,
		    TFTransfoXLMainLayer,
		    TFTransfoXLModel,
		    TFTransfoXLPreTrainedModel,
		)

else:
	import sys

	_lowerCAmelCase  :Dict  				=    _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)




'''simple docstring'''





import gc
import tempfile
import unittest

import numpy as np
import torch

from diffusers import VersatileDiffusionTextToImagePipeline
from diffusers.utils.testing_utils import nightly, require_torch_gpu, torch_device


_lowerCAmelCase  :List[Any]  				=    False






class 				UpperCAmelCase (       unittest.TestCase ):




    '''simple docstring'''







    pass






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




    '''simple docstring'''







    def       _UpperCamelCase    (       self )							-> Any:
         # clean up the VRAM after each test
         super().tearDown()
         gc.collect()
         torch.cuda.empty_cache()






    def       _UpperCamelCase    (       self )							-> Tuple:
         SCREAMING_SNAKE_CASE					:			Dict    	=       VersatileDiffusionTextToImagePipeline.from_pretrained('shi-labs/versatile-diffusion' )
         # remove text_unet
         pipe.remove_unused_weights()
         pipe.to(lowercase__ )
         pipe.set_progress_bar_config(disable=lowercase__ )

         SCREAMING_SNAKE_CASE					:			Union[str, Any]    	=       'A painting of a squirrel eating a burger '
         SCREAMING_SNAKE_CASE					:			List[Any]    	=       torch.manual_seed(0 )
         SCREAMING_SNAKE_CASE					:			Union[str, Any]    	=       pipe(
             prompt=lowercase__    ,							generator=lowercase__    ,							guidance_scale=7.5    ,							num_inference_steps=2    ,							output_type='numpy' ).images

         with tempfile.TemporaryDirectory() as tmpdirname:
              pipe.save_pretrained(lowercase__ )
              SCREAMING_SNAKE_CASE					:			int    	=       VersatileDiffusionTextToImagePipeline.from_pretrained(lowercase__ )
         pipe.to(lowercase__ )
         pipe.set_progress_bar_config(disable=lowercase__ )

         SCREAMING_SNAKE_CASE					:			List[str]    	=       generator.manual_seed(0 )
         SCREAMING_SNAKE_CASE					:			List[str]    	=       pipe(
             prompt=lowercase__    ,							generator=lowercase__    ,							guidance_scale=7.5    ,							num_inference_steps=2    ,							output_type='numpy' ).images

         assert np.abs(image - new_image ).sum() < 1E-5, "Models don't have the same forward pass"






    def       _UpperCamelCase    (       self )							-> Any:
         SCREAMING_SNAKE_CASE					:			Optional[Any]    	=       VersatileDiffusionTextToImagePipeline.from_pretrained(
             'shi-labs/versatile-diffusion'    ,							torch_dtype=torch.floataa )
         pipe.to(lowercase__ )
         pipe.set_progress_bar_config(disable=lowercase__ )

         SCREAMING_SNAKE_CASE					:			int    	=       'A painting of a squirrel eating a burger '
         SCREAMING_SNAKE_CASE					:			Tuple    	=       torch.manual_seed(0 )
         SCREAMING_SNAKE_CASE					:			Tuple    	=       pipe(
             prompt=lowercase__    ,							generator=lowercase__    ,							guidance_scale=7.5    ,							num_inference_steps=50    ,							output_type='numpy' ).images

         SCREAMING_SNAKE_CASE					:			Dict    	=       image[0, 253:256, 253:256, -1]

         assert image.shape == (1, 512, 512, 3)
         SCREAMING_SNAKE_CASE					:			Optional[Any]    	=       np.array([0.3_3_6_7, 0.3_1_6_9, 0.2_6_5_6, 0.3_8_7_0, 0.4_7_9_0, 0.3_7_9_6, 0.4_0_0_9, 0.4_8_7_8, 0.4_7_7_8] )

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



