






from manim import *


class 							__UpperCamelCase     (  _lowerCAmelCase ):







	def 			_a (	self   :				str      )		->			Tuple:
		"""simple docstring"""







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

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

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

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

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

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

		__lowercase   							=	[]
		__lowercase   							=	[]

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

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

		self.add(*_lowerCAmelCase							,					*_lowerCAmelCase      )

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

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

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

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

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

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

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

		self.play(Write(_lowerCAmelCase      )      )

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

		self.play(FadeOut(_lowerCAmelCase      )      )

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

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

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

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

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

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

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

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

			self.play(_lowerCAmelCase      )

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

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

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

			__lowercase   							=	a_c
			__lowercase   							=	a_c.copy()

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

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

		self.wait()


'''simple docstring'''


from typing import TYPE_CHECKING

# rely on isort to merge the imports
from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


__snake_case						  =		{"""configuration_focalnet""": ["""FOCALNET_PRETRAINED_CONFIG_ARCHIVE_MAP""", """FocalNetConfig"""]}


try:
       if not is_torch_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       __snake_case						  =		[
           """FOCALNET_PRETRAINED_MODEL_ARCHIVE_LIST""",
           """FocalNetForImageClassification""",
           """FocalNetForMaskedImageModeling""",
           """FocalNetBackbone""",
           """FocalNetModel""",
           """FocalNetPreTrainedModel""",
       ]

if TYPE_CHECKING:
       from .configuration_focalnet import FOCALNET_PRETRAINED_CONFIG_ARCHIVE_MAP, FocalNetConfig

       try:
              if not is_torch_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_focalnet import (
                  FOCALNET_PRETRAINED_MODEL_ARCHIVE_LIST,
                  FocalNetBackbone,
                  FocalNetForImageClassification,
                  FocalNetForMaskedImageModeling,
                  FocalNetModel,
                  FocalNetPreTrainedModel,
              )

else:
       import sys

       __snake_case						  =		_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)

