


"""simple docstring"""







import gc
import unittest

from diffusers import FlaxDPMSolverMultistepScheduler, FlaxStableDiffusionPipeline
from diffusers.utils import is_flax_available, slow
from diffusers.utils.testing_utils import require_flax


if is_flax_available():
		import jax
		import jax.numpy as jnp
		from flax.jax_utils import replicate
		from flax.training.common_utils import shard




@slow
@require_flax
class  UpperCamelCase_	( unittest.TestCase):

	"""simple docstring"""





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




	def   UpperCAmelCase_      ( self			:				Dict   )       ->	Dict:
						__SCREAMING_SNAKE_CASE   ,  __SCREAMING_SNAKE_CASE						   =				FlaxStableDiffusionPipeline.from_pretrained(
						    "stabilityai/stable-diffusion-2"			,			revision="bf16"			,			dtype=jnp.bfloataa			,			)

						__SCREAMING_SNAKE_CASE						   =				"A painting of a squirrel eating a burger"
						__SCREAMING_SNAKE_CASE						   =				jax.device_count()
						__SCREAMING_SNAKE_CASE						   =				num_samples * [prompt]
						__SCREAMING_SNAKE_CASE						   =				sd_pipe.prepare_inputs(UpperCAmelCase__   )

						__SCREAMING_SNAKE_CASE						   =				replicate(UpperCAmelCase__   )
						__SCREAMING_SNAKE_CASE						   =				shard(UpperCAmelCase__   )

						__SCREAMING_SNAKE_CASE						   =				jax.random.PRNGKey(0   )
						__SCREAMING_SNAKE_CASE						   =				jax.random.split(UpperCAmelCase__			,			jax.device_count()   )

						__SCREAMING_SNAKE_CASE						   =				sd_pipe(UpperCAmelCase__			,			UpperCAmelCase__			,			UpperCAmelCase__			,			num_inference_steps=2_5			,			jit=UpperCAmelCase__   )[0]
						assert images.shape == (jax.device_count(), 1, 7_6_8, 7_6_8, 3)

						__SCREAMING_SNAKE_CASE						   =				images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]   )
						__SCREAMING_SNAKE_CASE						   =				images[0, 2_5_3:2_5_6, 2_5_3:2_5_6, -1]

						__SCREAMING_SNAKE_CASE						   =				jnp.asarray(jax.device_get(image_slice.flatten()   )   )
						__SCREAMING_SNAKE_CASE						   =				jnp.array([0.4_238, 0.4_414, 0.4_395, 0.4_453, 0.4_629, 0.4_590, 0.4_531, 0.45_508, 0.4_512]   )
						print(F"""output_slice: {output_slice}"""   )
						assert jnp.abs(output_slice - expected_slice   ).max() < 1E-2




	def   UpperCAmelCase_      ( self			:				Union[str, Any]   )       ->	List[Any]:
						__SCREAMING_SNAKE_CASE						   =				"stabilityai/stable-diffusion-2"
						__SCREAMING_SNAKE_CASE   ,  __SCREAMING_SNAKE_CASE						   =				FlaxDPMSolverMultistepScheduler.from_pretrained(UpperCAmelCase__			,			subfolder="scheduler"   )
						__SCREAMING_SNAKE_CASE   ,  __SCREAMING_SNAKE_CASE						   =				FlaxStableDiffusionPipeline.from_pretrained(
						    UpperCAmelCase__			,			scheduler=UpperCAmelCase__			,			revision="bf16"			,			dtype=jnp.bfloataa			,			)
						__SCREAMING_SNAKE_CASE						   =				scheduler_params

						__SCREAMING_SNAKE_CASE						   =				"A painting of a squirrel eating a burger"
						__SCREAMING_SNAKE_CASE						   =				jax.device_count()
						__SCREAMING_SNAKE_CASE						   =				num_samples * [prompt]
						__SCREAMING_SNAKE_CASE						   =				sd_pipe.prepare_inputs(UpperCAmelCase__   )

						__SCREAMING_SNAKE_CASE						   =				replicate(UpperCAmelCase__   )
						__SCREAMING_SNAKE_CASE						   =				shard(UpperCAmelCase__   )

						__SCREAMING_SNAKE_CASE						   =				jax.random.PRNGKey(0   )
						__SCREAMING_SNAKE_CASE						   =				jax.random.split(UpperCAmelCase__			,			jax.device_count()   )

						__SCREAMING_SNAKE_CASE						   =				sd_pipe(UpperCAmelCase__			,			UpperCAmelCase__			,			UpperCAmelCase__			,			num_inference_steps=2_5			,			jit=UpperCAmelCase__   )[0]
						assert images.shape == (jax.device_count(), 1, 7_6_8, 7_6_8, 3)

						__SCREAMING_SNAKE_CASE						   =				images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]   )
						__SCREAMING_SNAKE_CASE						   =				images[0, 2_5_3:2_5_6, 2_5_3:2_5_6, -1]

						__SCREAMING_SNAKE_CASE						   =				jnp.asarray(jax.device_get(image_slice.flatten()   )   )
						__SCREAMING_SNAKE_CASE						   =				jnp.array([0.4_336, 0.42_969, 0.4_453, 0.4_199, 0.4_297, 0.4_531, 0.4_434, 0.4_434, 0.4_297]   )
						print(F"""output_slice: {output_slice}"""   )
						assert jnp.abs(output_slice - expected_slice   ).max() < 1E-2





from typing import TYPE_CHECKING

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


__UpperCAmelCase          =       {
    '''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:
    __UpperCAmelCase          =       [
        '''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:
    __UpperCAmelCase          =       [
        '''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

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



