



from manim import *

class  __snake_case      (   SCREAMING_SNAKE_CASE):







							'''simple docstring'''






							def _a      (							self		):
											a__				      =		Rectangle(height=0.5  ,  width=0.5		)
											a__				      =		Rectangle(height=0.25  ,  width=0.25		)
											a__				      =		Rectangle(height=0.46  ,  width=0.46		).set_stroke(width=0		)

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

											a__				      =		[mem.copy() for i in range(4		)]
											a__				      =		VGroup(*a_		).arrange(a_  ,  buff=0		)
											a__				      =		Text("""GPU"""  ,  font_size=24		)
											a__				      =		Group(a_  ,  a_		).arrange(a_  ,  buff=0.5  ,  aligned_edge=a_		)
											gpu.move_to([-1, -1, 0]		)
											self.add(a_		)

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

											a__				      =		Text("""Model"""  ,  font_size=24		)
											a__				      =		Group(a_  ,  a_		).arrange(a_  ,  buff=0.5  ,  aligned_edge=a_		)
											model.move_to([3, -1.0, 0]		)
											self.add(a_		)

											a__				      =		[]
											a__				      =		[]
											a__				      =		[]

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

															a__				      =		Rectangle(height=0.46 / 4  ,  width=0.46 / 3		).set_stroke(width=0.0		).set_fill(a_  ,  opacity=0.7		)

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

											self.add(*a_  ,  *a_  ,  *a_		)

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

											a__				      =		Text("""Loaded Checkpoint"""  ,  font_size=24		)
											a__				      =		Group(a_  ,  a_		).arrange(a_  ,  buff=0.5  ,  aligned_edge=a_		)
											checkpoint.move_to([3, 0.5, 0]		)

											self.add(a_		)

											a__				      =		[]
											a__				      =		[]

											for i, rect in enumerate(a_		):
															a__				      =		fill.copy().set_fill(a_  ,  opacity=0.7		)
															target.move_to(a_		)
															ckpt_arr.append(a_		)

															a__				      =		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(a_		)
											self.add(*a_  ,  *a_		)

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

											a__				      =		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(a_  ,  a_		)

											a__				      =		MarkupText(
											    F'''<span fgcolor=\'{BLUE}\'>●</span> Checkpoint'''  ,  font_size=18  ,  )

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

											a__				      =		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]		)

											a__				      =		[meta_mem.copy() for i in range(6		)]
											a__				      =		[meta_mem.copy() for i in range(6		)]
											a__				      =		VGroup(*a_		).arrange(a_  ,  buff=0		)
											a__				      =		VGroup(*a_		).arrange(a_  ,  buff=0		)
											a__				      =		VGroup(a_  ,  a_		).arrange(a_  ,  buff=0		)
											a__				      =		Text("""Disk"""  ,  font_size=24		)
											a__				      =		Group(a_  ,  a_		).arrange(a_  ,  buff=0.5  ,  aligned_edge=a_		)
											disk.move_to([-4.0, -1.25, 0]		)
											self.play(Write(a_  ,  run_time=3		)  ,  Write(a_  ,  run_time=1		)  ,  Create(a_  ,  run_time=1		)		)

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

											self.play(FadeOut(a_		)		)

											a__				      =		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(a_  ,  run_time=3		)		)

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

											self.wait()






import gc
import unittest

import numpy as np
import torch

from diffusers import (
    AudioDiffusionPipeline,
    AutoencoderKL,
    DDIMScheduler,
    DDPMScheduler,
    DiffusionPipeline,
    Mel,
    UNetaDConditionModel,
    UNetaDModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()

class  __snake_case      (   unittest.TestCase):







							'''simple docstring'''






							def _a      (							self		):
											# clean up the VRAM after each test
											super().tearDown()
											gc.collect()
											torch.cuda.empty_cache()


							@property
							def _a      (							self		):
											torch.manual_seed(0		)
											a__				      =		UNetaDModel(
											    sample_size=(32, 64)  ,  in_channels=1  ,  out_channels=1  ,  layers_per_block=2  ,  block_out_channels=(128, 128)  ,  down_block_types=("""AttnDownBlock2D""", """DownBlock2D""")  ,  up_block_types=("""UpBlock2D""", """AttnUpBlock2D""")  ,  )
											return model


							@property
							def _a      (							self		):
											torch.manual_seed(0		)
											a__				      =		UNetaDConditionModel(
											    sample_size=(64, 32)  ,  in_channels=1  ,  out_channels=1  ,  layers_per_block=2  ,  block_out_channels=(128, 128)  ,  down_block_types=("""CrossAttnDownBlock2D""", """DownBlock2D""")  ,  up_block_types=("""UpBlock2D""", """CrossAttnUpBlock2D""")  ,  cross_attention_dim=10  ,  )
											return model


							@property
							def _a      (							self		):
											torch.manual_seed(0		)
											a__				      =		AutoencoderKL(
											    sample_size=(128, 64)  ,  in_channels=1  ,  out_channels=1  ,  latent_channels=1  ,  layers_per_block=2  ,  block_out_channels=(128, 128)  ,  down_block_types=("""DownEncoderBlock2D""", """DownEncoderBlock2D""")  ,  up_block_types=("""UpDecoderBlock2D""", """UpDecoderBlock2D""")  ,  )
											a__				      =		UNetaDModel(
											    sample_size=(64, 32)  ,  in_channels=1  ,  out_channels=1  ,  layers_per_block=2  ,  block_out_channels=(128, 128)  ,  down_block_types=("""AttnDownBlock2D""", """DownBlock2D""")  ,  up_block_types=("""UpBlock2D""", """AttnUpBlock2D""")  ,  )
											return vqvae, unet







							@slow
							def _a      (							self		):
											a__				      =		"""cpu"""  # ensure determinism for the device-dependent torch.Generator
											a__				      =		Mel(
											    x_res=self.dummy_unet.config.sample_size[1]  ,  y_res=self.dummy_unet.config.sample_size[0]  ,  )

											a__				      =		DDPMScheduler()
											a__				      =		AudioDiffusionPipeline(vqvae=a_  ,  unet=self.dummy_unet  ,  mel=a_  ,  scheduler=a_		)
											a__				      =		pipe.to(a_		)
											pipe.set_progress_bar_config(disable=a_		)

											a__				      =		torch.Generator(device=a_		).manual_seed(42		)
											a__				      =		pipe(generator=a_  ,  steps=4		)
											a__				      =		output.audios[0]
											a__				      =		output.images[0]

											a__				      =		torch.Generator(device=a_		).manual_seed(42		)
											a__				      =		pipe(generator=a_  ,  steps=4  ,  return_dict=a_		)
											a__				      =		output[0][0]

											assert audio.shape == (1, (self.dummy_unet.config.sample_size[1] - 1) * mel.hop_length)
											assert (
											    image.height == self.dummy_unet.config.sample_size[0]
											    and image.width == self.dummy_unet.config.sample_size[1]
											)
											a__				      =		np.frombuffer(image.tobytes()  ,  dtype="""uint8"""		)[:10]
											a__				      =		np.frombuffer(image_from_tuple.tobytes()  ,  dtype="""uint8"""		)[:10]
											a__				      =		np.array([69, 255, 255, 255, 0, 0, 77, 181, 12, 127]		)

											assert np.abs(image_slice.flatten() - expected_slice		).max() == 0
											assert np.abs(image_from_tuple_slice.flatten() - expected_slice		).max() == 0

											a__				      =		Mel(
											    x_res=self.dummy_vqvae_and_unet[0].config.sample_size[1]  ,  y_res=self.dummy_vqvae_and_unet[0].config.sample_size[0]  ,  )

											a__				      =		DDIMScheduler()
											a__				      =		self.dummy_vqvae_and_unet
											a__				      =		AudioDiffusionPipeline(
											    vqvae=self.dummy_vqvae_and_unet[0]  ,  unet=dummy_vqvae_and_unet[1]  ,  mel=a_  ,  scheduler=a_		)
											a__				      =		pipe.to(a_		)
											pipe.set_progress_bar_config(disable=a_		)

											np.random.seed(0		)
											a__				      =		np.random.uniform(-1  ,  1  ,  ((dummy_vqvae_and_unet[0].config.sample_size[1] - 1) * mel.hop_length,)		)
											a__				      =		torch.Generator(device=a_		).manual_seed(42		)
											a__				      =		pipe(raw_audio=a_  ,  generator=a_  ,  start_step=5  ,  steps=10		)
											a__				      =		output.images[0]

											assert (
											    image.height == self.dummy_vqvae_and_unet[0].config.sample_size[0]
											    and image.width == self.dummy_vqvae_and_unet[0].config.sample_size[1]
											)
											a__				      =		np.frombuffer(image.tobytes()  ,  dtype="""uint8"""		)[:10]
											a__				      =		np.array([120, 117, 110, 109, 138, 167, 138, 148, 132, 121]		)

											assert np.abs(image_slice.flatten() - expected_slice		).max() == 0

											a__				      =		self.dummy_unet_condition
											a__				      =		AudioDiffusionPipeline(
											    vqvae=self.dummy_vqvae_and_unet[0]  ,  unet=a_  ,  mel=a_  ,  scheduler=a_		)
											a__				      =		pipe.to(a_		)
											pipe.set_progress_bar_config(disable=a_		)

											np.random.seed(0		)
											a__				      =		torch.rand((1, 1, 10)		)
											a__				      =		pipe(generator=a_  ,  encoding=a_		)
											a__				      =		output.images[0]
											a__				      =		np.frombuffer(image.tobytes()  ,  dtype="""uint8"""		)[:10]
											a__				      =		np.array([107, 103, 120, 127, 142, 122, 113, 122, 97, 111]		)

											assert np.abs(image_slice.flatten() - expected_slice		).max() == 0




@slow
@require_torch_gpu
class  __snake_case      (   unittest.TestCase):







							'''simple docstring'''






							def _a      (							self		):
											# clean up the VRAM after each test
											super().tearDown()
											gc.collect()
											torch.cuda.empty_cache()







							def _a      (							self		):
											a__				      =		torch_device

											a__				      =		DiffusionPipeline.from_pretrained("""teticio/audio-diffusion-ddim-256"""		)
											a__				      =		pipe.to(a_		)
											pipe.set_progress_bar_config(disable=a_		)

											a__				      =		torch.Generator(device=a_		).manual_seed(42		)
											a__				      =		pipe(generator=a_		)
											a__				      =		output.audios[0]
											a__				      =		output.images[0]

											assert audio.shape == (1, (pipe.unet.config.sample_size[1] - 1) * pipe.mel.hop_length)
											assert image.height == pipe.unet.config.sample_size[0] and image.width == pipe.unet.config.sample_size[1]
											a__				      =		np.frombuffer(image.tobytes()  ,  dtype="""uint8"""		)[:10]
											a__				      =		np.array([151, 167, 154, 144, 122, 134, 121, 105, 70, 26]		)

											assert np.abs(image_slice.flatten() - expected_slice		).max() == 0

