





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    ):





     def 		UpperCamelCase(     self					):
        # clean up the VRAM after each test
        super().tearDown()
        gc.collect()





     def 		UpperCamelCase(     self					):
        _UpperCAmelCase       ,      _UpperCAmelCase           =   FlaxStableDiffusionPipeline.from_pretrained(
            '''stabilityai/stable-diffusion-2'''				,							revision='''bf16'''				,							dtype=jnp.bfloataa				,							)

        _UpperCAmelCase           =   '''A painting of a squirrel eating a burger'''
        _UpperCAmelCase           =   jax.device_count()
        _UpperCAmelCase           =   num_samples * [prompt]
        _UpperCAmelCase           =   sd_pipe.prepare_inputs(_UpperCamelCase					)

        _UpperCAmelCase           =   replicate(_UpperCamelCase					)
        _UpperCAmelCase           =   shard(_UpperCamelCase					)

        _UpperCAmelCase           =   jax.random.PRNGKey(0					)
        _UpperCAmelCase           =   jax.random.split(_UpperCamelCase				,							jax.device_count()					)

        _UpperCAmelCase           =   sd_pipe(_UpperCamelCase				,							_UpperCamelCase				,							_UpperCamelCase				,							num_inference_steps=25				,							jit=_UpperCamelCase					)[0]
        assert images.shape == (jax.device_count(), 1, 768, 768, 3)

        _UpperCAmelCase           =   images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]					)
        _UpperCAmelCase           =   images[0, 253:256, 253:256, -1]

        _UpperCAmelCase           =   jnp.asarray(jax.device_get(image_slice.flatten()					)					)
        _UpperCAmelCase           =   jnp.array([0.4238, 0.4414, 0.4395, 0.4453, 0.4629, 0.4590, 0.4531, 0.45508, 0.4512]					)
        print(f'''output_slice: {output_slice}'''					)
        assert jnp.abs(output_slice - expected_slice					).max() < 1e-2





     def 		UpperCamelCase(     self					):
        _UpperCAmelCase           =   '''stabilityai/stable-diffusion-2'''
        _UpperCAmelCase       ,      _UpperCAmelCase           =   FlaxDPMSolverMultistepScheduler.from_pretrained(_UpperCamelCase				,							subfolder='''scheduler'''					)
        _UpperCAmelCase       ,      _UpperCAmelCase           =   FlaxStableDiffusionPipeline.from_pretrained(
            _UpperCamelCase				,							scheduler=_UpperCamelCase				,							revision='''bf16'''				,							dtype=jnp.bfloataa				,							)
        _UpperCAmelCase           =   scheduler_params

        _UpperCAmelCase           =   '''A painting of a squirrel eating a burger'''
        _UpperCAmelCase           =   jax.device_count()
        _UpperCAmelCase           =   num_samples * [prompt]
        _UpperCAmelCase           =   sd_pipe.prepare_inputs(_UpperCamelCase					)

        _UpperCAmelCase           =   replicate(_UpperCamelCase					)
        _UpperCAmelCase           =   shard(_UpperCamelCase					)

        _UpperCAmelCase           =   jax.random.PRNGKey(0					)
        _UpperCAmelCase           =   jax.random.split(_UpperCamelCase				,							jax.device_count()					)

        _UpperCAmelCase           =   sd_pipe(_UpperCamelCase				,							_UpperCamelCase				,							_UpperCamelCase				,							num_inference_steps=25				,							jit=_UpperCamelCase					)[0]
        assert images.shape == (jax.device_count(), 1, 768, 768, 3)

        _UpperCAmelCase           =   images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]					)
        _UpperCAmelCase           =   images[0, 253:256, 253:256, -1]

        _UpperCAmelCase           =   jnp.asarray(jax.device_get(image_slice.flatten()					)					)
        _UpperCAmelCase           =   jnp.array([0.4336, 0.42969, 0.4453, 0.4199, 0.4297, 0.4531, 0.4434, 0.4434, 0.4297]					)
        print(f'''output_slice: {output_slice}'''					)
        assert jnp.abs(output_slice - expected_slice					).max() < 1e-2






from dataclasses import dataclass, field
from typing import ClassVar, Dict

from ..features import Features, Value
from .base import TaskTemplate




@dataclass(frozen=A__    )
class 			__UpperCamelCase    (							A__    ):
     __A	:  str							  =   field(default="""language-modeling"""	, metadata={"""include_in_asdict_even_if_is_default""": True}    )

     __A	:  ClassVar[Features]							  =   Features({"""text""": Value("""string"""    )}    )
     __A	:  ClassVar[Features]							  =   Features({}    )
     __A	:  str							  =   "text"





     @property
     def 		UpperCamelCase(     self					):
        return {self.text_column: "text"}