




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



 '''simple docstring'''


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

       """simple docstring"""






       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()

 @property
 def        a__				(       self      :						str						)     ->			Any:

       """simple docstring"""






       torch.manual_seed(0						)
       lowerCamelCase_	  =					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      :						Optional[Any]						)     ->			Union[str, Any]:

       """simple docstring"""






       torch.manual_seed(0						)
       lowerCamelCase_	  =					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      :						Tuple						)     ->			int:

       """simple docstring"""






       torch.manual_seed(0						)
       lowerCamelCase_	  =					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')		,							)
       lowerCamelCase_	  =					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      :						Dict						)     ->			Optional[int]:

       """simple docstring"""






       lowerCamelCase_	  =					'cpu'  # ensure determinism for the device-dependent torch.Generator
       lowerCamelCase_	  =					Mel(
           x_res=self.dummy_unet.config.sample_size[1]		,							y_res=self.dummy_unet.config.sample_size[0]		,							)

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

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

       lowerCamelCase_	  =					torch.Generator(device=A_						).manual_seed(42						)
       lowerCamelCase_	  =					pipe(generator=A_		,							steps=4		,							return_dict=A_						)
       lowerCamelCase_	  =					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]
       )
       lowerCamelCase_	  =					np.frombuffer(image.tobytes()		,							dtype='uint8'						)[:10]
       lowerCamelCase_	  =					np.frombuffer(image_from_tuple.tobytes()		,							dtype='uint8'						)[:10]
       lowerCamelCase_	  =					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

       lowerCamelCase_	  =					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]		,							)

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

       np.random.seed(0						)
       lowerCamelCase_	  =					np.random.uniform(-1		,							1		,							((dummy_vqvae_and_unet[0].config.sample_size[1] - 1) * mel.hop_length,)						)
       lowerCamelCase_	  =					torch.Generator(device=A_						).manual_seed(42						)
       lowerCamelCase_	  =					pipe(raw_audio=A_		,							generator=A_		,							start_step=5		,							steps=10						)
       lowerCamelCase_	  =					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]
       )
       lowerCamelCase_	  =					np.frombuffer(image.tobytes()		,							dtype='uint8'						)[:10]
       lowerCamelCase_	  =					np.array([120, 117, 110, 109, 138, 167, 138, 148, 132, 121]						)

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

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

       np.random.seed(0						)
       lowerCamelCase_	  =					torch.rand((1, 1, 10)						)
       lowerCamelCase_	  =					pipe(generator=A_		,							encoding=A_						)
       lowerCamelCase_	  =					output.images[0]
       lowerCamelCase_	  =					np.frombuffer(image.tobytes()		,							dtype='uint8'						)[:10]
       lowerCamelCase_	  =					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 						A(		unittest.TestCase     ):



 '''simple docstring'''


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

       """simple docstring"""






       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()





 def        a__				(       self      :						str						)     ->			Any:

       """simple docstring"""






       lowerCamelCase_	  =					torch_device

       lowerCamelCase_	  =					DiffusionPipeline.from_pretrained('teticio/audio-diffusion-ddim-256'						)
       lowerCamelCase_	  =					pipe.to(A_						)
       pipe.set_progress_bar_config(disable=A_						)

       lowerCamelCase_	  =					torch.Generator(device=A_						).manual_seed(42						)
       lowerCamelCase_	  =					pipe(generator=A_						)
       lowerCamelCase_	  =					output.audios[0]
       lowerCamelCase_	  =					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]
       lowerCamelCase_	  =					np.frombuffer(image.tobytes()		,							dtype='uint8'						)[:10]
       lowerCamelCase_	  =					np.array([151, 167, 154, 144, 122, 134, 121, 105, 70, 26]						)

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









from manim import *







class 						A(		UpperCamelCase     ):



 '''simple docstring'''


 def        a__				(       self      :						Optional[Any]						)     ->			List[str]:

       """simple docstring"""






       lowerCamelCase_	  =					Rectangle(height=0.5		,							width=0.5						)
       lowerCamelCase_	  =					Rectangle(height=0.46		,							width=0.46						).set_stroke(width=0						)

       lowerCamelCase_	  =					Rectangle(height=0.25		,							width=0.25						)

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

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

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

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

       lowerCamelCase_	  =					[]
       lowerCamelCase_	  =					[]

       for i, rect in enumerate(A_						):
             lowerCamelCase_	  =					fill.copy().set_fill(A_		,							opacity=0.8						)
             target.move_to(A_						)
             model_arr.append(A_						)

             lowerCamelCase_	  =					Rectangle(height=0.46		,							width=0.46						).set_stroke(width=0.0						).set_fill(A_		,							opacity=0.8						)
             cpu_target.move_to(cpu_left_col_base[i]						)
             model_cpu_arr.append(A_						)

       self.add(*A_		,							*A_						)

       lowerCamelCase_	  =					[meta_mem.copy() for i in range(6						)]
       lowerCamelCase_	  =					[meta_mem.copy() for i in range(6						)]
       lowerCamelCase_	  =					VGroup(*A_						).arrange(A_		,							buff=0						)
       lowerCamelCase_	  =					VGroup(*A_						).arrange(A_		,							buff=0						)
       lowerCamelCase_	  =					VGroup(A_		,							A_						).arrange(A_		,							buff=0						)
       lowerCamelCase_	  =					Text('Disk'		,							font_size=24						)
       lowerCamelCase_	  =					Group(A_		,							A_						).arrange(A_		,							buff=0.5		,							aligned_edge=A_						)
       disk.move_to([-4, -1.25, 0]						)
       self.add(A_		,							A_						)

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

       self.add(A_		,							A_						)

       lowerCamelCase_	  =					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_						)

       lowerCamelCase_	  =					MarkupText(
           f"""Now watch as an input is passed through the model\nand how the memory is utilized and handled."""		,							font_size=24		,							)
       step_a.move_to([2, 2, 0]						)

       self.play(Write(A_						)						)

       lowerCamelCase_	  =					Square(0.3						)
       input.set_fill(A_		,							opacity=1.0						)
       input.set_stroke(width=0.0						)
       input.next_to(model_base[0]		,							A_		,							buff=0.5						)

       self.play(Write(A_						)						)

       input.generate_target()
       input.target.next_to(model_arr[0]		,							direction=A_		,							buff=0.02						)
       self.play(MoveToTarget(A_						)						)

       self.play(FadeOut(A_						)						)

       lowerCamelCase_	  =					Arrow(start=A_		,							end=A_		,							color=A_		,							buff=0.5						)
       a.next_to(model_arr[0].get_left()		,							A_		,							buff=0.2						)

       model_cpu_arr[0].generate_target()
       model_cpu_arr[0].target.move_to(gpu_rect[0]						)

       lowerCamelCase_	  =					MarkupText(
           f"""As the input reaches a layer, the hook triggers\nand weights are moved from the CPU\nto the GPU and back."""		,							font_size=24		,							)
       step_a.move_to([2, 2, 0]						)

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

       lowerCamelCase_	  =					{'run_time': 1, 'fade_in': True, 'fade_out': True, 'buff': 0.02}

       self.play(
           Write(A_						)		,							Circumscribe(model_arr[0]		,							color=A_		,							**A_						)		,							Circumscribe(model_cpu_arr[0]		,							color=A_		,							**A_						)		,							Circumscribe(gpu_rect[0]		,							color=A_		,							**A_						)		,							)
       self.play(MoveToTarget(model_cpu_arr[0]						)						)

       lowerCamelCase_	  =					a.copy()
       for i in range(6						):
             a_c.next_to(model_arr[i].get_right() + 0.02		,							A_		,							buff=0.2						)

             input.generate_target()
             input.target.move_to(model_arr[i].get_right() + 0.02						)

             lowerCamelCase_	  =					AnimationGroup(
                 FadeOut(A_		,							run_time=0.5						)		,							MoveToTarget(A_		,							run_time=0.5						)		,							FadeIn(A_		,							run_time=0.5						)		,							lag_ratio=0.2						)

             self.play(A_						)

             model_cpu_arr[i].generate_target()
             model_cpu_arr[i].target.move_to(cpu_left_col_base[i]						)

             if i < 5:
                   model_cpu_arr[i + 1].generate_target()
                   model_cpu_arr[i + 1].target.move_to(gpu_rect[0]						)
                   if i >= 1:
                         lowerCamelCase_	  =					0.7

                   self.play(
                       Circumscribe(model_arr[i]		,							**A_						)		,							Circumscribe(cpu_left_col_base[i]		,							**A_						)		,							Circumscribe(cpu_left_col_base[i + 1]		,							color=A_		,							**A_						)		,							Circumscribe(gpu_rect[0]		,							color=A_		,							**A_						)		,							Circumscribe(model_arr[i + 1]		,							color=A_		,							**A_						)		,							)
                   if i < 1:
                         self.play(
                             MoveToTarget(model_cpu_arr[i]						)		,							MoveToTarget(model_cpu_arr[i + 1]						)		,							)
                   else:
                         self.play(
                             MoveToTarget(model_cpu_arr[i]		,							run_time=0.7						)		,							MoveToTarget(model_cpu_arr[i + 1]		,							run_time=0.7						)		,							)
             else:
                   model_cpu_arr[i].generate_target()
                   model_cpu_arr[i].target.move_to(cpu_left_col_base[-1]						)
                   input.generate_target()
                   input.target.next_to(model_arr[-1].get_right()		,							RIGHT + 0.02		,							buff=0.2						)

                   self.play(
                       Circumscribe(model_arr[-1]		,							color=A_		,							**A_						)		,							Circumscribe(cpu_left_col_base[-1]		,							color=A_		,							**A_						)		,							Circumscribe(gpu_rect[0]		,							color=A_		,							**A_						)		,							)

                   self.play(MoveToTarget(model_cpu_arr[i]						)						)

             lowerCamelCase_	  =					a_c
             lowerCamelCase_	  =					a_c.copy()

       input.generate_target()
       input.target.next_to(model_base[-1]		,							RIGHT + 0.02		,							buff=0.5						)
       self.play(
           FadeOut(A_						)		,							FadeOut(A_		,							run_time=0.5						)		,							)

       lowerCamelCase_	  =					MarkupText(f"""Inference on a model too large for GPU memory\nis successfully completed."""		,							font_size=24						)
       step_a.move_to([2, 2, 0]						)

       self.play(Write(A_		,							run_time=3						)		,							MoveToTarget(A_						)						)

       self.wait()



