





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

   '''simple docstring'''


   def    UpperCamelCase__						(					self							)				->   List[Any]:






       """simple docstring"""
       super().tearDown()
       gc.collect()




   def    UpperCamelCase__						(					self							)				->   Optional[Any]:






       """simple docstring"""
       UpperCamelCase__							:       Any      =   FlaxStableDiffusionPipeline.from_pretrained(
           '''stabilityai/stable-diffusion-2''',	revision='''bf16''',	dtype=jnp.bfloataa,	)

       UpperCamelCase__							:       List[str]      =   """A painting of a squirrel eating a burger"""
       UpperCamelCase__							:       Any      =   jax.device_count()
       UpperCamelCase__							:       List[str]      =   num_samples * [prompt]
       UpperCamelCase__							:       Optional[int]      =   sd_pipe.prepare_inputs(_lowercase							)

       UpperCamelCase__							:       Optional[Any]      =   replicate(_lowercase							)
       UpperCamelCase__							:       Any      =   shard(_lowercase							)

       UpperCamelCase__							:       List[Any]      =   jax.random.PRNGKey(0							)
       UpperCamelCase__							:       Dict      =   jax.random.split(_lowercase,	jax.device_count()							)

       UpperCamelCase__							:       List[str]      =   sd_pipe(_lowercase,	_lowercase,	_lowercase,	num_inference_steps=25,	jit=_lowercase							)[0]
       assert images.shape == (jax.device_count(), 1, 768, 768, 3)

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

       UpperCamelCase__							:       Union[str, Any]      =   jnp.asarray(jax.device_get(image_slice.flatten()							)							)
       UpperCamelCase__							:       str      =   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							)				->   Optional[Any]:






       """simple docstring"""
       UpperCamelCase__							:       str      =   """stabilityai/stable-diffusion-2"""
       UpperCamelCase__							:       int      =   FlaxDPMSolverMultistepScheduler.from_pretrained(_lowercase,	subfolder='''scheduler'''							)
       UpperCamelCase__							:       Any      =   FlaxStableDiffusionPipeline.from_pretrained(
           _lowercase,	scheduler=_lowercase,	revision='''bf16''',	dtype=jnp.bfloataa,	)
       UpperCamelCase__							:       List[str]      =   scheduler_params

       UpperCamelCase__							:       List[str]      =   """A painting of a squirrel eating a burger"""
       UpperCamelCase__							:       List[str]      =   jax.device_count()
       UpperCamelCase__							:       Dict      =   num_samples * [prompt]
       UpperCamelCase__							:       str      =   sd_pipe.prepare_inputs(_lowercase							)

       UpperCamelCase__							:       Dict      =   replicate(_lowercase							)
       UpperCamelCase__							:       Tuple      =   shard(_lowercase							)

       UpperCamelCase__							:       Tuple      =   jax.random.PRNGKey(0							)
       UpperCamelCase__							:       Dict      =   jax.random.split(_lowercase,	jax.device_count()							)

       UpperCamelCase__							:       Union[str, Any]      =   sd_pipe(_lowercase,	_lowercase,	_lowercase,	num_inference_steps=25,	jit=_lowercase							)[0]
       assert images.shape == (jax.device_count(), 1, 768, 768, 3)

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

       UpperCamelCase__							:       List[Any]      =   jnp.asarray(jax.device_get(image_slice.flatten()							)							)
       UpperCamelCase__							:       List[str]      =   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








from manim import *






class 					lowercase__ (			__lowerCamelCase       ):

   '''simple docstring'''


   def    UpperCamelCase__						(					self							)				->   List[str]:






       """simple docstring"""
       UpperCamelCase__							:       Any      =   Rectangle(height=0.5,	width=0.5							)
       UpperCamelCase__							:       Any      =   Rectangle(height=0.25,	width=0.25							)
       UpperCamelCase__							:       Dict      =   Rectangle(height=0.46,	width=0.46							).set_stroke(width=0							)

       UpperCamelCase__							:       List[Any]      =   [mem.copy() for i in range(6							)]
       UpperCamelCase__							:       Optional[int]      =   [mem.copy() for i in range(6							)]
       UpperCamelCase__							:       str      =   VGroup(*__magic_name__							).arrange(__magic_name__,	buff=0							)
       UpperCamelCase__							:       List[str]      =   VGroup(*__magic_name__							).arrange(__magic_name__,	buff=0							)
       UpperCamelCase__							:       Tuple      =   VGroup(__magic_name__,	__magic_name__							).arrange(__magic_name__,	buff=0							)
       UpperCamelCase__							:       List[str]      =   Text('''CPU''',	font_size=24							)
       UpperCamelCase__							:       Any      =   Group(__magic_name__,	__magic_name__							).arrange(__magic_name__,	buff=0.5,	aligned_edge=__magic_name__							)
       cpu.move_to([-2.5, -0.5, 0]							)
       self.add(__magic_name__							)

       UpperCamelCase__							:       int      =   [mem.copy() for i in range(4							)]
       UpperCamelCase__							:       str      =   VGroup(*__magic_name__							).arrange(__magic_name__,	buff=0							)
       UpperCamelCase__							:       Tuple      =   Text('''GPU''',	font_size=24							)
       UpperCamelCase__							:       str      =   Group(__magic_name__,	__magic_name__							).arrange(__magic_name__,	buff=0.5,	aligned_edge=__magic_name__							)
       gpu.move_to([-1, -1, 0]							)
       self.add(__magic_name__							)

       UpperCamelCase__							:       str      =   [mem.copy() for i in range(6							)]
       UpperCamelCase__							:       Tuple      =   VGroup(*__magic_name__							).arrange(__magic_name__,	buff=0							)

       UpperCamelCase__							:       Tuple      =   Text('''Model''',	font_size=24							)
       UpperCamelCase__							:       Tuple      =   Group(__magic_name__,	__magic_name__							).arrange(__magic_name__,	buff=0.5,	aligned_edge=__magic_name__							)
       model.move_to([3, -1.0, 0]							)
       self.add(__magic_name__							)

       UpperCamelCase__							:       Union[str, Any]      =   []
       UpperCamelCase__							:       Any      =   []
       UpperCamelCase__							:       List[Any]      =   []

       for i, rect in enumerate(__magic_name__							):
           rect.set_stroke(__magic_name__							)

           UpperCamelCase__							:       Tuple      =   Rectangle(height=0.46 / 4,	width=0.46 / 3							).set_stroke(width=0.0							).set_fill(__magic_name__,	opacity=0.7							)

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

       self.add(*__magic_name__,	*__magic_name__,	*__magic_name__							)

       UpperCamelCase__							:       Tuple      =   [mem.copy() for i in range(6							)]
       UpperCamelCase__							:       Dict      =   VGroup(*__magic_name__							).arrange(__magic_name__,	buff=0							)

       UpperCamelCase__							:       List[str]      =   Text('''Loaded Checkpoint''',	font_size=24							)
       UpperCamelCase__							:       Optional[Any]      =   Group(__magic_name__,	__magic_name__							).arrange(__magic_name__,	buff=0.5,	aligned_edge=__magic_name__							)
       checkpoint.move_to([3, 0.5, 0]							)

       self.add(__magic_name__							)

       UpperCamelCase__							:       Tuple      =   []
       UpperCamelCase__							:       List[str]      =   []

       for i, rect in enumerate(__magic_name__							):
           UpperCamelCase__							:       Optional[int]      =   fill.copy().set_fill(__magic_name__,	opacity=0.7							)
           target.move_to(__magic_name__							)
           ckpt_arr.append(__magic_name__							)

           UpperCamelCase__							:       int      =   target.copy()
           if i < 5:
               cpu_target.move_to(cpu_left_col_base[i + 1]							)
           else:
               cpu_target.move_to(cpu_right_col_base[i - 5]							)
           ckpt_cpu_arr.append(__magic_name__							)
       self.add(*__magic_name__,	*__magic_name__							)

       UpperCamelCase__							:       List[str]      =   Square(side_length=2.2							)
       key.move_to([-5, 2, 0]							)

       UpperCamelCase__							:       List[str]      =   MarkupText(
           f"<b>Key:</b>\n\n<span fgcolor='{YELLOW}'>●</span> Empty Model",	font_size=18,	)

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

       self.add(__magic_name__,	__magic_name__							)

       UpperCamelCase__							:       Any      =   MarkupText(
           f"<span fgcolor='{BLUE}'>●</span> Checkpoint",	font_size=18,	)

       blue_text.next_to(__magic_name__,	DOWN * 2.4,	aligned_edge=key_text.get_left()							)
       self.add(__magic_name__							)

       UpperCamelCase__							:       Dict      =   MarkupText(
           f"Based on the passed in configuration, weights are stored in\na variety of np.memmaps on disk or to a particular device.",	font_size=24,	)
       step_a.move_to([2, 2, 0]							)

       UpperCamelCase__							:       Any      =   [meta_mem.copy() for i in range(6							)]
       UpperCamelCase__							:       int      =   [meta_mem.copy() for i in range(6							)]
       UpperCamelCase__							:       Optional[int]      =   VGroup(*__magic_name__							).arrange(__magic_name__,	buff=0							)
       UpperCamelCase__							:       Union[str, Any]      =   VGroup(*__magic_name__							).arrange(__magic_name__,	buff=0							)
       UpperCamelCase__							:       Union[str, Any]      =   VGroup(__magic_name__,	__magic_name__							).arrange(__magic_name__,	buff=0							)
       UpperCamelCase__							:       Any      =   Text('''Disk''',	font_size=24							)
       UpperCamelCase__							:       List[Any]      =   Group(__magic_name__,	__magic_name__							).arrange(__magic_name__,	buff=0.5,	aligned_edge=__magic_name__							)
       disk.move_to([-4.0, -1.25, 0]							)
       self.play(Write(__magic_name__,	run_time=3							),	Write(__magic_name__,	run_time=1							),	Create(__magic_name__,	run_time=1							)							)

       UpperCamelCase__							:       Union[str, Any]      =   []
       for i, rect in enumerate(__magic_name__							):
           UpperCamelCase__							:       Union[str, Any]      =   rect.copy()
           target.generate_target()
           target.target.move_to(disk_left_col_base[i]							).scale(0.5							)
           animations.append(MoveToTarget(__magic_name__,	run_time=1.5							)							)
       self.play(*__magic_name__							)

       self.play(FadeOut(__magic_name__							)							)

       UpperCamelCase__							:       Optional[int]      =   MarkupText(f"Then, the checkpoint is removed from memory\nthrough garbage collection.",	font_size=24							)
       step_a.move_to([2, 2, 0]							)

       self.play(Write(__magic_name__,	run_time=3							)							)

       self.play(
           FadeOut(__magic_name__,	__magic_name__,	*__magic_name__,	*__magic_name__							),	)

       self.wait()

