




from manim import *

class        lowercase      (    __lowerCamelCase):







				def 					a_  (							self   :					Union[str, Any]       ):



					"""simple docstring"""

					A_  :				str	   =		Rectangle(height=0.5      ,     width=0.5       )
					A_  :				str	   =		Rectangle(height=0.46      ,     width=0.46       ).set_stroke(width=0       )

					A_  :				int	   =		Rectangle(height=0.25      ,     width=0.25       )

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

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

					A_  :				Optional[int]	   =		[mem.copy() for i in range(6       )]
					A_  :				Optional[Any]	   =		VGroup(*UpperCamelCase_       ).arrange(UpperCamelCase_      ,     buff=0       )

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

					A_  :				Union[str, Any]	   =		[]
					A_  :				Union[str, Any]	   =		[]

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

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

					self.add(*UpperCamelCase_      ,     *UpperCamelCase_       )

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

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

					A_  :				Optional[int]	   =		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(UpperCamelCase_      ,     UpperCamelCase_       )

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

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

					A_  :				Union[str, Any]	   =		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(UpperCamelCase_       )       )

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

					self.play(Write(UpperCamelCase_       )       )

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

					self.play(FadeOut(UpperCamelCase_       )       )

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

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

					A_  :				Optional[Any]	   =		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(UpperCamelCase_      ,     run_time=3       )       )

					A_  :				Union[str, Any]	   =		{'''run_time''': 1, '''fade_in''': True, '''fade_out''': True, '''buff''': 0.02}

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

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

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

						A_  :				Dict	   =		AnimationGroup(
						    FadeOut(UpperCamelCase_      ,     run_time=0.5       )      ,     MoveToTarget(UpperCamelCase_      ,     run_time=0.5       )      ,     FadeIn(UpperCamelCase_      ,     run_time=0.5       )      ,     lag_ratio=0.2       )

						self.play(UpperCamelCase_       )

						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:
								A_  :				List[Any]	   =		0.7

							self.play(
							    Circumscribe(model_arr[i]      ,     **UpperCamelCase_       )      ,     Circumscribe(cpu_left_col_base[i]      ,     **UpperCamelCase_       )      ,     Circumscribe(cpu_left_col_base[i + 1]      ,     color=UpperCamelCase_      ,     **UpperCamelCase_       )      ,     Circumscribe(gpu_rect[0]      ,     color=UpperCamelCase_      ,     **UpperCamelCase_       )      ,     Circumscribe(model_arr[i + 1]      ,     color=UpperCamelCase_      ,     **UpperCamelCase_       )      ,     )
							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=UpperCamelCase_      ,     **UpperCamelCase_       )      ,     Circumscribe(cpu_left_col_base[-1]      ,     color=UpperCamelCase_      ,     **UpperCamelCase_       )      ,     Circumscribe(gpu_rect[0]      ,     color=UpperCamelCase_      ,     **UpperCamelCase_       )      ,     )

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

						A_  :				Optional[Any]	   =		a_c
						A_  :				List[str]	   =		a_c.copy()

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

					A_  :				Optional[int]	   =		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(UpperCamelCase_      ,     run_time=3       )      ,     MoveToTarget(UpperCamelCase_       )       )

					self.wait()











"""simple docstring"""

import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import TransformeraDModel, VQDiffusionPipeline, VQDiffusionScheduler, VQModel
from diffusers.pipelines.vq_diffusion.pipeline_vq_diffusion import LearnedClassifierFreeSamplingEmbeddings
from diffusers.utils import load_numpy, slow, torch_device
from diffusers.utils.testing_utils import require_torch_gpu


_lowerCamelCase  :			Any        =       False

class        lowercase      (    unittest.TestCase):







				def 					a_  (							self   :					Tuple       ):



					"""simple docstring"""

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







				@property
				def 					a_  (							self   :					Any       ):



					"""simple docstring"""

					return 12







				@property
				def 					a_  (							self   :					List[str]       ):



					"""simple docstring"""

					return 12







				@property
				def 					a_  (							self   :					List[Any]       ):



					"""simple docstring"""

					return 32







				@property
				def 					a_  (							self   :					Any       ):



					"""simple docstring"""

					torch.manual_seed(0       )
					A_  :				Union[str, Any]	   =		VQModel(
					    block_out_channels=[32, 64]      ,     in_channels=3      ,     out_channels=3      ,     down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']      ,     up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']      ,     latent_channels=3      ,     num_vq_embeddings=self.num_embed      ,     vq_embed_dim=3      ,     )
					return model







				@property
				def 					a_  (							self   :					List[Any]       ):



					"""simple docstring"""

					A_  :				int	   =		CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''       )
					return tokenizer







				@property
				def 					a_  (							self   :					int       ):



					"""simple docstring"""

					torch.manual_seed(0       )
					A_  :				Optional[int]	   =		CLIPTextConfig(
					    bos_token_id=0      ,     eos_token_id=2      ,     hidden_size=self.text_embedder_hidden_size      ,     intermediate_size=37      ,     layer_norm_eps=1E-05      ,     num_attention_heads=4      ,     num_hidden_layers=5      ,     pad_token_id=1      ,     vocab_size=10_00      ,     )
					return CLIPTextModel(_lowerCamelCase       )







				@property
				def 					a_  (							self   :					Optional[int]       ):



					"""simple docstring"""

					torch.manual_seed(0       )

					A_  :				Optional[Any]	   =		12
					A_  :				Optional[int]	   =		12

					A_  :				int	   =		{
					    '''attention_bias''': True,
					    '''cross_attention_dim''': 32,
					    '''attention_head_dim''': height * width,
					    '''num_attention_heads''': 1,
					    '''num_vector_embeds''': self.num_embed,
					    '''num_embeds_ada_norm''': self.num_embeds_ada_norm,
					    '''norm_num_groups''': 32,
					    '''sample_size''': width,
					    '''activation_fn''': '''geglu-approximate''',
					}

					A_  :				Tuple	   =		TransformeraDModel(**_lowerCamelCase       )
					return model







				def 					a_  (							self   :					Optional[int]       ):



					"""simple docstring"""

					A_  :				Union[str, Any]	   =		'''cpu'''

					A_  :				Union[str, Any]	   =		self.dummy_vqvae
					A_  :				str	   =		self.dummy_text_encoder
					A_  :				List[Any]	   =		self.dummy_tokenizer
					A_  :				int	   =		self.dummy_transformer
					A_  :				Any	   =		VQDiffusionScheduler(self.num_embed       )
					A_  :				Optional[Any]	   =		LearnedClassifierFreeSamplingEmbeddings(learnable=_lowerCamelCase       )

					A_  :				Dict	   =		VQDiffusionPipeline(
					    vqvae=_lowerCamelCase      ,     text_encoder=_lowerCamelCase      ,     tokenizer=_lowerCamelCase      ,     transformer=_lowerCamelCase      ,     scheduler=_lowerCamelCase      ,     learned_classifier_free_sampling_embeddings=_lowerCamelCase      ,     )
					A_  :				List[Any]	   =		pipe.to(_lowerCamelCase       )
					pipe.set_progress_bar_config(disable=_lowerCamelCase       )

					A_  :				List[Any]	   =		'''teddy bear playing in the pool'''

					A_  :				List[Any]	   =		torch.Generator(device=_lowerCamelCase       ).manual_seed(0       )
					A_  :				List[Any]	   =		pipe([prompt]      ,     generator=_lowerCamelCase      ,     num_inference_steps=2      ,     output_type='''np'''       )
					A_  :				Any	   =		output.images

					A_  :				List[str]	   =		torch.Generator(device=_lowerCamelCase       ).manual_seed(0       )
					A_  :				List[Any]	   =		pipe(
					    [prompt]      ,     generator=_lowerCamelCase      ,     output_type='''np'''      ,     return_dict=_lowerCamelCase      ,     num_inference_steps=2       )[0]

					A_  :				Optional[int]	   =		image[0, -3:, -3:, -1]
					A_  :				Optional[int]	   =		image_from_tuple[0, -3:, -3:, -1]

					assert image.shape == (1, 24, 24, 3)

					A_  :				Optional[int]	   =		np.array([0.6551, 0.6168, 0.5008, 0.5676, 0.5659, 0.4295, 0.6073, 0.5599, 0.4992]       )

					assert np.abs(image_slice.flatten() - expected_slice       ).max() < 1E-2
					assert np.abs(image_from_tuple_slice.flatten() - expected_slice       ).max() < 1E-2







				def 					a_  (							self   :					List[Any]       ):



					"""simple docstring"""

					A_  :				Union[str, Any]	   =		'''cpu'''

					A_  :				int	   =		self.dummy_vqvae
					A_  :				List[str]	   =		self.dummy_text_encoder
					A_  :				Optional[Any]	   =		self.dummy_tokenizer
					A_  :				Any	   =		self.dummy_transformer
					A_  :				Any	   =		VQDiffusionScheduler(self.num_embed       )
					A_  :				Optional[int]	   =		LearnedClassifierFreeSamplingEmbeddings(
					    learnable=_lowerCamelCase      ,     hidden_size=self.text_embedder_hidden_size      ,     length=tokenizer.model_max_length       )

					A_  :				int	   =		VQDiffusionPipeline(
					    vqvae=_lowerCamelCase      ,     text_encoder=_lowerCamelCase      ,     tokenizer=_lowerCamelCase      ,     transformer=_lowerCamelCase      ,     scheduler=_lowerCamelCase      ,     learned_classifier_free_sampling_embeddings=_lowerCamelCase      ,     )
					A_  :				List[Any]	   =		pipe.to(_lowerCamelCase       )
					pipe.set_progress_bar_config(disable=_lowerCamelCase       )

					A_  :				Any	   =		'''teddy bear playing in the pool'''

					A_  :				str	   =		torch.Generator(device=_lowerCamelCase       ).manual_seed(0       )
					A_  :				Optional[Any]	   =		pipe([prompt]      ,     generator=_lowerCamelCase      ,     num_inference_steps=2      ,     output_type='''np'''       )
					A_  :				Tuple	   =		output.images

					A_  :				Optional[int]	   =		torch.Generator(device=_lowerCamelCase       ).manual_seed(0       )
					A_  :				List[str]	   =		pipe(
					    [prompt]      ,     generator=_lowerCamelCase      ,     output_type='''np'''      ,     return_dict=_lowerCamelCase      ,     num_inference_steps=2       )[0]

					A_  :				Optional[int]	   =		image[0, -3:, -3:, -1]
					A_  :				Optional[Any]	   =		image_from_tuple[0, -3:, -3:, -1]

					assert image.shape == (1, 24, 24, 3)

					A_  :				str	   =		np.array([0.6693, 0.6075, 0.4959, 0.5701, 0.5583, 0.4333, 0.6171, 0.5684, 0.4988]       )

					assert np.abs(image_slice.flatten() - expected_slice       ).max() < 2.0
					assert np.abs(image_from_tuple_slice.flatten() - expected_slice       ).max() < 1E-2

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







				def 					a_  (							self   :					Any       ):



					"""simple docstring"""

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







				def 					a_  (							self   :					str       ):



					"""simple docstring"""

					A_  :				str	   =		load_numpy(
					    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
					    '''/vq_diffusion/teddy_bear_pool_classifier_free_sampling.npy'''       )

					A_  :				int	   =		VQDiffusionPipeline.from_pretrained('''microsoft/vq-diffusion-ithq'''       )
					A_  :				Tuple	   =		pipeline.to(_lowerCamelCase       )
					pipeline.set_progress_bar_config(disable=_lowerCamelCase       )

					# requires GPU generator for gumbel softmax
					# don't use GPU generator in tests though
					A_  :				Dict	   =		torch.Generator(device=_lowerCamelCase       ).manual_seed(0       )
					A_  :				Union[str, Any]	   =		pipeline(
					    '''teddy bear playing in the pool'''      ,     num_images_per_prompt=1      ,     generator=_lowerCamelCase      ,     output_type='''np'''      ,     )

					A_  :				Optional[int]	   =		output.images[0]

					assert image.shape == (2_56, 2_56, 3)
					assert np.abs(expected_image - image       ).max() < 2.0





