



from manim import *






class    SCREAMING_SNAKE_CASE			(							lowerCamelCase__					):





							def 						UpperCamelCase_					(   self		:		List[Any]	):

											'''simple docstring'''
											__a       =		Rectangle(height=0.5      ,				width=0.5	)
											__a       =		Rectangle(height=0.46      ,				width=0.46	).set_stroke(width=0	)

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

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

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

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

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

											__a       =		[]
											__a       =		[]

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

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

											self.add(*__lowercase      ,				*__lowercase	)

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

											__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(__lowercase      ,				__lowercase	)

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

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

											__a       =		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(__lowercase	)	)

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

											self.play(Write(__lowercase	)	)

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

											self.play(FadeOut(__lowercase	)	)

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

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

											__a       =		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(__lowercase      ,				run_time=3	)	)

											__a       =		{"""run_time""": 1, """fade_in""": True, """fade_out""": True, """buff""": 0.02}

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

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

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

															__a       =		AnimationGroup(
															    FadeOut(__lowercase      ,				run_time=0.5	)      ,				MoveToTarget(__lowercase      ,				run_time=0.5	)      ,				FadeIn(__lowercase      ,				run_time=0.5	)      ,				lag_ratio=0.2	)

															self.play(__lowercase	)

															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       =		0.7

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

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

															__a       =		a_c
															__a       =		a_c.copy()

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

											__a       =		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(__lowercase      ,				run_time=3	)      ,				MoveToTarget(__lowercase	)	)

											self.wait()





from typing import Dict, List, Optional, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import center_crop, normalize, rescale, resize, to_channel_dimension_format
from ...image_utils import (
    IMAGENET_STANDARD_MEAN,
    IMAGENET_STANDARD_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_vision_available, logging


if is_vision_available():
						import PIL


lowerCamelCase__			= logging.get_logger(__name__)






class    SCREAMING_SNAKE_CASE			(							lowerCamelCase__					):
							__lowerCamelCase	:	str    =['pixel_values']





							def __init__(   self		:		str      ,				__lowercase		:		bool = True      ,				__lowercase		:		Dict[str, int] = None      ,				__lowercase		:		PILImageResampling = PIL.Image.BICUBIC      ,				__lowercase		:		bool = True      ,				__lowercase		:		Dict[str, int] = None      ,				__lowercase		:		Union[int, float] = 1 / 255      ,				__lowercase		:		bool = True      ,				__lowercase		:		bool = True      ,				__lowercase		:		Optional[Union[float, List[float]]] = None      ,				__lowercase		:		Optional[Union[float, List[float]]] = None      ,				**__lowercase		:		Any      ,				):

											'''simple docstring'''
											super().__init__(**__lowercase	)
											__a       =		size if size is not None else {"""height""": 256, """width""": 256}
											__a       =		get_size_dict(__lowercase	)
											__a       =		crop_size if crop_size is not None else {"""height""": 224, """width""": 224}
											__a       =		get_size_dict(__lowercase      ,				param_name="""crop_size"""	)

											__a       =		do_resize
											__a       =		size
											__a       =		resample
											__a       =		do_center_crop
											__a       =		crop_size
											__a       =		do_rescale
											__a       =		rescale_factor
											__a       =		do_normalize
											__a       =		image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN
											__a       =		image_std if image_std is not None else IMAGENET_STANDARD_STD





							def 						UpperCamelCase_					(   self		:		List[Any]      ,				__lowercase		:		np.ndarray      ,				__lowercase		:		Dict[str, int]      ,				__lowercase		:		PILImageResampling = PIL.Image.BICUBIC      ,				__lowercase		:		Optional[Union[str, ChannelDimension]] = None      ,				**__lowercase		:		Union[str, Any]      ,				):

											'''simple docstring'''
											__a       =		get_size_dict(__lowercase	)
											if "height" not in size or "width" not in size:
															raise ValueError(F"The size dictionary must have keys 'height' and 'width'. Got {size.keys()}"	)
											return resize(
											    __lowercase      ,				size=(size["""height"""], size["""width"""])      ,				resample=__lowercase      ,				data_format=__lowercase      ,				**__lowercase	)





							def 						UpperCamelCase_					(   self		:		Dict      ,				__lowercase		:		np.ndarray      ,				__lowercase		:		Dict[str, int]      ,				__lowercase		:		Optional[Union[str, ChannelDimension]] = None      ,				**__lowercase		:		str      ,				):

											'''simple docstring'''
											__a       =		get_size_dict(__lowercase	)
											if "height" not in size or "width" not in size:
															raise ValueError(F"The size dictionary must have keys 'height' and 'width'. Got {size.keys()}"	)
											return center_crop(__lowercase      ,				size=(size["""height"""], size["""width"""])      ,				data_format=__lowercase      ,				**__lowercase	)





							def 						UpperCamelCase_					(   self		:		int      ,				__lowercase		:		np.ndarray      ,				__lowercase		:		Union[int, float]      ,				__lowercase		:		Optional[Union[str, ChannelDimension]] = None      ,				**__lowercase		:		List[str]      ,				):

											'''simple docstring'''
											return rescale(__lowercase      ,				scale=__lowercase      ,				data_format=__lowercase      ,				**__lowercase	)





							def 						UpperCamelCase_					(   self		:		Tuple      ,				__lowercase		:		np.ndarray      ,				__lowercase		:		Union[float, List[float]]      ,				__lowercase		:		Union[float, List[float]]      ,				__lowercase		:		Optional[Union[str, ChannelDimension]] = None      ,				**__lowercase		:		str      ,				):

											'''simple docstring'''
											return normalize(__lowercase      ,				mean=__lowercase      ,				std=__lowercase      ,				data_format=__lowercase      ,				**__lowercase	)







							def 						UpperCamelCase_					(   self		:		Optional[Any]      ,				__lowercase		:		ImageInput      ,				__lowercase		:		bool = None      ,				__lowercase		:		Dict[str, int] = None      ,				__lowercase		:		Dict=None      ,				__lowercase		:		bool = None      ,				__lowercase		:		Dict[str, int] = None      ,				__lowercase		:		bool = None      ,				__lowercase		:		float = None      ,				__lowercase		:		bool = None      ,				__lowercase		:		Optional[Union[float, List[float]]] = None      ,				__lowercase		:		Optional[Union[float, List[float]]] = None      ,				__lowercase		:		Optional[Union[str, TensorType]] = None      ,				__lowercase		:		ChannelDimension = ChannelDimension.FIRST      ,				**__lowercase		:		Optional[Any]      ,				):

											'''simple docstring'''
											__a       =		do_resize if do_resize is not None else self.do_resize
											__a       =		resample if resample is not None else self.resample
											__a       =		do_center_crop if do_center_crop is not None else self.do_center_crop
											__a       =		do_rescale if do_rescale is not None else self.do_rescale
											__a       =		rescale_factor if rescale_factor is not None else self.rescale_factor
											__a       =		do_normalize if do_normalize is not None else self.do_normalize
											__a       =		image_mean if image_mean is not None else self.image_mean
											__a       =		image_std if image_std is not None else self.image_std

											__a       =		size if size is not None else self.size
											__a       =		get_size_dict(__lowercase	)
											__a       =		crop_size if crop_size is not None else self.crop_size
											__a       =		get_size_dict(__lowercase      ,				param_name="""crop_size"""	)

											__a       =		make_list_of_images(__lowercase	)

											if not valid_images(__lowercase	):
															raise ValueError(
															    """Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, """
															    """torch.Tensor, tf.Tensor or jax.ndarray."""	)

											if do_resize and size is None or resample is None:
															raise ValueError("""Size and resample must be specified if do_resize is True."""	)

											if do_center_crop and crop_size is None:
															raise ValueError("""Crop size must be specified if do_center_crop is True."""	)

											if do_rescale and rescale_factor is None:
															raise ValueError("""Rescale factor must be specified if do_rescale is True."""	)

											if do_normalize and (image_mean is None or image_std is None):
															raise ValueError("""Image mean and std must be specified if do_normalize is True."""	)

											# All transformations expect numpy arrays.
											__a       =		[to_numpy_array(__lowercase	) for image in images]

											if do_resize:
															__a       =		[self.resize(image=__lowercase      ,				size=__lowercase      ,				resample=__lowercase	) for image in images]

											if do_center_crop:
															__a       =		[self.center_crop(image=__lowercase      ,				size=__lowercase	) for image in images]

											if do_rescale:
															__a       =		[self.rescale(image=__lowercase      ,				scale=__lowercase	) for image in images]

											if do_normalize:
															__a       =		[self.normalize(image=__lowercase      ,				mean=__lowercase      ,				std=__lowercase	) for image in images]

											__a       =		[to_channel_dimension_format(__lowercase      ,				__lowercase	) for image in images]

											__a       =		{"""pixel_values""": images}
											return BatchFeature(data=__lowercase      ,				tensor_type=__lowercase	)
