






'''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       ):
								# clean up the VRAM after each test
								super().tearDown()
								gc.collect()






		def 			UpperCamelCase_					(      self       ):
								lowerCamelCase__       ,			lowerCamelCase__         =  FlaxStableDiffusionPipeline.from_pretrained(
								    """stabilityai/stable-diffusion-2"""   ,revision="""bf16"""   ,dtype=jnp.bfloataa   ,)

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

								lowerCamelCase__         =  replicate(_lowerCAmelCase       )
								lowerCamelCase__         =  shard(_lowerCAmelCase       )

								lowerCamelCase__         =  jax.random.PRNGKey(0       )
								lowerCamelCase__         =  jax.random.split(_lowerCAmelCase   ,jax.device_count()       )

								lowerCamelCase__         =  sd_pipe(_lowerCAmelCase   ,_lowerCAmelCase   ,_lowerCAmelCase   ,num_inference_steps=25   ,jit=_lowerCAmelCase       )[0]
								assert images.shape == (jax.device_count(), 1, 7_68, 7_68, 3)

								lowerCamelCase__         =  images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]       )
								lowerCamelCase__         =  images[0, 2_53:2_56, 2_53:2_56, -1]

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






		def 			UpperCamelCase_					(      self       ):
								lowerCamelCase__         =  """stabilityai/stable-diffusion-2"""
								lowerCamelCase__       ,			lowerCamelCase__         =  FlaxDPMSolverMultistepScheduler.from_pretrained(_lowerCAmelCase   ,subfolder="""scheduler"""       )
								lowerCamelCase__       ,			lowerCamelCase__         =  FlaxStableDiffusionPipeline.from_pretrained(
								    _lowerCAmelCase   ,scheduler=_lowerCAmelCase   ,revision="""bf16"""   ,dtype=jnp.bfloataa   ,)
								lowerCamelCase__         =  scheduler_params

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

								lowerCamelCase__         =  replicate(_lowerCAmelCase       )
								lowerCamelCase__         =  shard(_lowerCAmelCase       )

								lowerCamelCase__         =  jax.random.PRNGKey(0       )
								lowerCamelCase__         =  jax.random.split(_lowerCAmelCase   ,jax.device_count()       )

								lowerCamelCase__         =  sd_pipe(_lowerCAmelCase   ,_lowerCAmelCase   ,_lowerCAmelCase   ,num_inference_steps=25   ,jit=_lowerCAmelCase       )[0]
								assert images.shape == (jax.device_count(), 1, 7_68, 7_68, 3)

								lowerCamelCase__         =  images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]       )
								lowerCamelCase__         =  images[0, 2_53:2_56, 2_53:2_56, -1]

								lowerCamelCase__         =  jnp.asarray(jax.device_get(image_slice.flatten()       )       )
								lowerCamelCase__         =  jnp.array([0.4336, 0.4_2969, 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









'''simple docstring'''







from manim import *




class 		UpperCamelCase__    (a			):




		'''simple docstring'''





		def 			UpperCamelCase_					(      self       ):
								lowerCamelCase__         =  Rectangle(height=0.5   ,width=0.5       )
								lowerCamelCase__         =  Rectangle(height=0.46   ,width=0.46       ).set_stroke(width=0       )

								lowerCamelCase__         =  [mem.copy() for i in range(6       )]
								lowerCamelCase__         =  [mem.copy() for i in range(6       )]
								lowerCamelCase__         =  VGroup(*_lowerCAmelCase       ).arrange(_lowerCAmelCase   ,buff=0       )
								lowerCamelCase__         =  VGroup(*_lowerCAmelCase       ).arrange(_lowerCAmelCase   ,buff=0       )
								lowerCamelCase__         =  VGroup(_lowerCAmelCase   ,_lowerCAmelCase       ).arrange(_lowerCAmelCase   ,buff=0       )
								lowerCamelCase__         =  Text("""CPU"""   ,font_size=24       )
								lowerCamelCase__         =  Group(_lowerCAmelCase   ,_lowerCAmelCase       ).arrange(_lowerCAmelCase   ,buff=0.5   ,aligned_edge=_lowerCAmelCase       )
								cpu.move_to([-2.5, -0.5, 0]       )
								self.add(_lowerCAmelCase       )

								lowerCamelCase__         =  [mem.copy() for i in range(1       )]
								lowerCamelCase__         =  VGroup(*_lowerCAmelCase       ).arrange(_lowerCAmelCase   ,buff=0       )
								lowerCamelCase__         =  Text("""GPU"""   ,font_size=24       )
								lowerCamelCase__         =  Group(_lowerCAmelCase   ,_lowerCAmelCase       ).arrange(_lowerCAmelCase   ,buff=0.5   ,aligned_edge=_lowerCAmelCase       )
								gpu.align_to(_lowerCAmelCase   ,_lowerCAmelCase       )
								gpu.set_x(gpu.get_x() - 1       )

								self.add(_lowerCAmelCase       )

								lowerCamelCase__         =  [mem.copy() for i in range(6       )]
								lowerCamelCase__         =  VGroup(*_lowerCAmelCase       ).arrange(_lowerCAmelCase   ,buff=0       )

								lowerCamelCase__         =  Text("""Model"""   ,font_size=24       )
								lowerCamelCase__         =  Group(_lowerCAmelCase   ,_lowerCAmelCase       ).arrange(_lowerCAmelCase   ,buff=0.5   ,aligned_edge=_lowerCAmelCase       )
								model.move_to([3, -1.0, 0]       )

								self.play(
								    Create(_lowerCAmelCase   ,run_time=1       )   ,Create(_lowerCAmelCase   ,run_time=1       )   ,Create(_lowerCAmelCase   ,run_time=1       )   ,)

								lowerCamelCase__         =  MarkupText(
								    F'''First, an empty model skeleton is loaded\ninto <span fgcolor=\'{YELLOW}\'>memory</span> without using much RAM.'''   ,font_size=24   ,)

								lowerCamelCase__         =  Square(side_length=2.2       )
								key.move_to([-5, 2, 0]       )

								lowerCamelCase__         =  MarkupText(
								    F'''<b>Key:</b>\n\n<span fgcolor=\'{YELLOW}\'>●</span> Empty Model'''   ,font_size=18   ,)

								key_text.move_to([-5, 2.4, 0]       )

								step_a.move_to([2, 2, 0]       )
								self.play(Write(_lowerCAmelCase   ,run_time=2.5       )   ,Write(_lowerCAmelCase       )   ,Write(_lowerCAmelCase       )       )

								self.add(_lowerCAmelCase       )

								lowerCamelCase__         =  []
								lowerCamelCase__         =  []
								lowerCamelCase__         =  []
								for i, rect in enumerate(_lowerCAmelCase       ):
														lowerCamelCase__         =  Rectangle(height=0.46   ,width=0.46       ).set_stroke(width=0.0       ).set_fill(_lowerCAmelCase   ,opacity=0.7       )
														cpu_target.move_to(_lowerCAmelCase       )
														cpu_target.generate_target()
														lowerCamelCase__         =  0.46 / 4
														lowerCamelCase__         =  0.46 / 3

														if i == 0:
																				cpu_target.target.next_to(cpu_left_col_base[0].get_corner(DOWN + LEFT       )   ,buff=0.02   ,direction=_lowerCAmelCase       )
																				cpu_target.target.set_x(cpu_target.target.get_x() + 0.1       )
														elif i == 3:
																				cpu_target.target.next_to(cpu_targs[0].target   ,direction=_lowerCAmelCase   ,buff=0.0       )
														else:
																				cpu_target.target.next_to(cpu_targs[i - 1].target   ,direction=_lowerCAmelCase   ,buff=0.0       )
														cpu_targs.append(_lowerCAmelCase       )

														first_animations.append(rect.animate(run_time=0.5       ).set_stroke(_lowerCAmelCase       )       )
														second_animations.append(MoveToTarget(_lowerCAmelCase   ,run_time=1.5       )       )

								self.play(*_lowerCAmelCase       )
								self.play(*_lowerCAmelCase       )

								self.wait()

