






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     _a							(	unittest.TestCase     ):






  """simple docstring"""





  def  _UpperCAmelCase  (  self			)      ->			Union[str, Any]:
     # clean up the VRAM after each test
     super().tearDown()
     gc.collect()




  def  _UpperCAmelCase  (  self			)      ->			Optional[int]:
     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.4_2_3_8, 0.4_4_1_4, 0.4_3_9_5, 0.4_4_5_3, 0.4_6_2_9, 0.4_5_9_0, 0.4_5_3_1, 0.4_5_5_0_8, 0.4_5_1_2]			)
     print(f"""output_slice: {output_slice}"""			)
     assert jnp.abs(output_slice - expected_slice			).max() < 1e-2




  def  _UpperCAmelCase  (  self			)      ->			Union[str, Any]:
     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.4_3_3_6, 0.4_2_9_6_9, 0.4_4_5_3, 0.4_1_9_9, 0.4_2_9_7, 0.4_5_3_1, 0.4_4_3_4, 0.4_4_3_4, 0.4_2_9_7]			)
     print(f"""output_slice: {output_slice}"""			)
     assert jnp.abs(output_slice - expected_slice			).max() < 1e-2













import requests






def      _snake_case     (__lowercase     ,   __lowercase):
   UpperCamelCase_						=  {'Content-Type': 'application/json'}
   UpperCamelCase_						=  requests.post(__lowercase     ,   json={'text': message_body}     ,   headers=__lowercase)
   if response.status_code != 200:
      UpperCamelCase_						=  (
          'Request to slack returned an error '
          f"""{response.status_code}, the response is:\n{response.text}"""
      )
      raise ValueError(__lowercase)


if __name__ == "__main__":
 # Set the slack url to the one provided by Slack when you create the webhook at
 # https://my.slack.com/services/new/incoming-webhook/
 send_slack_message("""<YOUR MESSAGE BODY>""", """<SLACK CHANNEL URL>""")





