'''simple docstring'''

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,
    convert_to_rgb,
    get_resize_output_image_size,
    normalize,
    rescale,
    resize,
    to_channel_dimension_format,
)
from ...image_utils import (
    OPENAI_CLIP_MEAN,
    OPENAI_CLIP_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_vision_available, logging


__a:			int	       =       logging.get_logger(__name__)


if is_vision_available():
  import PIL



class       UpperCAmelCase  (						a__ ):





       '''simple docstring'''



       SCREAMING_SNAKE_CASE 		=   ["pixel_values"]

       def __init__(						self     ,       __lowerCAmelCase = True     ,       __lowerCAmelCase = None     ,       __lowerCAmelCase = PILImageResampling.BICUBIC     ,       __lowerCAmelCase = True     ,       __lowerCAmelCase = None     ,       __lowerCAmelCase = True     ,       __lowerCAmelCase = 1 / 255     ,       __lowerCAmelCase = True     ,       __lowerCAmelCase = None     ,       __lowerCAmelCase = None     ,       __lowerCAmelCase = True     ,       **__lowerCAmelCase     ,       )  ->  None:
        super().__init__(**__lowerCAmelCase			)
        lowercase__    :   int						= size if size is not None else {'''shortest_edge''': 224}
        lowercase__    :   Union[str, Any]						= get_size_dict(__lowerCAmelCase     ,       default_to_square=__lowerCAmelCase			)
        lowercase__    :   Union[str, Any]						= crop_size if crop_size is not None else {'''height''': 224, '''width''': 224}
        lowercase__    :   List[str]						= get_size_dict(__lowerCAmelCase     ,       default_to_square=__lowerCAmelCase     ,       param_name='''crop_size'''			)

        lowercase__    :   List[Any]						= do_resize
        lowercase__    :   Any						= size
        lowercase__    :   Dict						= resample
        lowercase__    :   List[Any]						= do_center_crop
        lowercase__    :   int						= crop_size
        lowercase__    :   Any						= do_rescale
        lowercase__    :   List[str]						= rescale_factor
        lowercase__    :   Tuple						= do_normalize
        lowercase__    :   int						= image_mean if image_mean is not None else OPENAI_CLIP_MEAN
        lowercase__    :   int						= image_std if image_std is not None else OPENAI_CLIP_STD
        lowercase__    :   Tuple						= do_convert_rgb

       def _lowerCAmelCase(						self     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase = PILImageResampling.BICUBIC     ,       __lowerCAmelCase = None     ,       **__lowerCAmelCase     ,       )  ->  np.ndarray:
        lowercase__    :   Union[str, Any]						= get_size_dict(__lowerCAmelCase     ,       default_to_square=__lowerCAmelCase			)
        if "shortest_edge" not in size:
         raise ValueError(F"""The `size` parameter must contain the key `shortest_edge`. Got {size.keys()}"""			)
        lowercase__    :   List[str]						= get_resize_output_image_size(__lowerCAmelCase     ,       size=size['''shortest_edge''']     ,       default_to_square=__lowerCAmelCase			)
        return resize(__lowerCAmelCase     ,       size=__lowerCAmelCase     ,       resample=__lowerCAmelCase     ,       data_format=__lowerCAmelCase     ,       **__lowerCAmelCase			)

       def _lowerCAmelCase(						self     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase = None     ,       **__lowerCAmelCase     ,       )  ->  np.ndarray:
        lowercase__    :   Dict						= get_size_dict(__lowerCAmelCase			)
        if "height" not in size or "width" not in size:
         raise ValueError(F"""The `size` parameter must contain the keys (height, width). Got {size.keys()}"""			)
        return center_crop(__lowerCAmelCase     ,       size=(size['''height'''], size['''width'''])     ,       data_format=__lowerCAmelCase     ,       **__lowerCAmelCase			)

       def _lowerCAmelCase(						self     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase = None     ,       **__lowerCAmelCase     ,       )  ->  Optional[int]:
        return rescale(__lowerCAmelCase     ,       scale=__lowerCAmelCase     ,       data_format=__lowerCAmelCase     ,       **__lowerCAmelCase			)

       def _lowerCAmelCase(						self     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase = None     ,       **__lowerCAmelCase     ,       )  ->  np.ndarray:
        return normalize(__lowerCAmelCase     ,       mean=__lowerCAmelCase     ,       std=__lowerCAmelCase     ,       data_format=__lowerCAmelCase     ,       **__lowerCAmelCase			)


       def _lowerCAmelCase(						self     ,       __lowerCAmelCase     ,       __lowerCAmelCase = None     ,       __lowerCAmelCase = None     ,       __lowerCAmelCase = None     ,       __lowerCAmelCase = None     ,       __lowerCAmelCase = None     ,       __lowerCAmelCase = None     ,       __lowerCAmelCase = None     ,       __lowerCAmelCase = None     ,       __lowerCAmelCase = None     ,       __lowerCAmelCase = None     ,       __lowerCAmelCase = None     ,       __lowerCAmelCase = None     ,       __lowerCAmelCase = ChannelDimension.FIRST     ,       **__lowerCAmelCase     ,       )  ->  PIL.Image.Image:
        lowercase__    :   Tuple						= do_resize if do_resize is not None else self.do_resize
        lowercase__    :   Optional[Any]						= size if size is not None else self.size
        lowercase__    :   Dict						= get_size_dict(__lowerCAmelCase     ,       param_name='''size'''     ,       default_to_square=__lowerCAmelCase			)
        lowercase__    :   Union[str, Any]						= resample if resample is not None else self.resample
        lowercase__    :   Union[str, Any]						= do_center_crop if do_center_crop is not None else self.do_center_crop
        lowercase__    :   Union[str, Any]						= crop_size if crop_size is not None else self.crop_size
        lowercase__    :   str						= get_size_dict(__lowerCAmelCase     ,       param_name='''crop_size'''     ,       default_to_square=__lowerCAmelCase			)
        lowercase__    :   Optional[Any]						= do_rescale if do_rescale is not None else self.do_rescale
        lowercase__    :   int						= rescale_factor if rescale_factor is not None else self.rescale_factor
        lowercase__    :   Optional[Any]						= do_normalize if do_normalize is not None else self.do_normalize
        lowercase__    :   str						= image_mean if image_mean is not None else self.image_mean
        lowercase__    :   int						= image_std if image_std is not None else self.image_std
        lowercase__    :   str						= do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb

        lowercase__    :   List[Any]						= make_list_of_images(__lowerCAmelCase			)

        if not valid_images(__lowerCAmelCase			):
         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:
         raise ValueError('''Size 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.'''			)

        # PIL RGBA images are converted to RGB
        if do_convert_rgb:
         lowercase__    :   List[str]						= [convert_to_rgb(__lowerCAmelCase			) for image in images]

        # All transformations expect numpy arrays.
        lowercase__    :   List[str]						= [to_numpy_array(__lowerCAmelCase			) for image in images]

        if do_resize:
         lowercase__    :   Any						= [self.resize(image=__lowerCAmelCase     ,       size=__lowerCAmelCase     ,       resample=__lowerCAmelCase			) for image in images]

        if do_center_crop:
         lowercase__    :   int						= [self.center_crop(image=__lowerCAmelCase     ,       size=__lowerCAmelCase			) for image in images]

        if do_rescale:
         lowercase__    :   Optional[Any]						= [self.rescale(image=__lowerCAmelCase     ,       scale=__lowerCAmelCase			) for image in images]

        if do_normalize:
         lowercase__    :   Tuple						= [self.normalize(image=__lowerCAmelCase     ,       mean=__lowerCAmelCase     ,       std=__lowerCAmelCase			) for image in images]

        lowercase__    :   List[str]						= [to_channel_dimension_format(__lowerCAmelCase     ,       __lowerCAmelCase			) for image in images]

        lowercase__    :   Dict						= {'''pixel_values''': images}
        return BatchFeature(data=__lowerCAmelCase     ,       tensor_type=__lowerCAmelCase			)


'''simple docstring'''

from manim import *



class       UpperCAmelCase  (						a__ ):





       '''simple docstring'''



       def _lowerCAmelCase(						self			)  ->  List[Any]:
        lowercase__    :   int						= Rectangle(height=0.5     ,       width=0.5			)
        lowercase__    :   Optional[int]						= Rectangle(height=0.2_5     ,       width=0.2_5			)
        lowercase__    :   Tuple						= Rectangle(height=0.4_6     ,       width=0.4_6			).set_stroke(width=0			)

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

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

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

        lowercase__    :   Any						= Text('''Model'''     ,       font_size=24			)
        lowercase__    :   Optional[Any]						= Group(__lowerCAmelCase     ,       __lowerCAmelCase			).arrange(__lowerCAmelCase     ,       buff=0.5     ,       aligned_edge=__lowerCAmelCase			)
        model.move_to([3, -1.0, 0]			)
        self.add(__lowerCAmelCase			)

        lowercase__    :   int						= []
        lowercase__    :   int						= []
        lowercase__    :   Any						= []

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

         lowercase__    :   Optional[int]						= Rectangle(height=0.4_6 / 4     ,       width=0.4_6 / 3			).set_stroke(width=0.0			).set_fill(__lowerCAmelCase     ,       opacity=0.7			)

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

        self.add(*__lowerCAmelCase     ,       *__lowerCAmelCase     ,       *__lowerCAmelCase			)

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

        lowercase__    :   Optional[Any]						= Text('''Loaded Checkpoint'''     ,       font_size=24			)
        lowercase__    :   Dict						= Group(__lowerCAmelCase     ,       __lowerCAmelCase			).arrange(__lowerCAmelCase     ,       buff=0.5     ,       aligned_edge=__lowerCAmelCase			)
        checkpoint.move_to([3, 0.5, 0]			)

        self.add(__lowerCAmelCase			)

        lowercase__    :   str						= []
        lowercase__    :   List[str]						= []

        for i, rect in enumerate(__lowerCAmelCase			):
         lowercase__    :   List[str]						= fill.copy().set_fill(__lowerCAmelCase     ,       opacity=0.7			)
         target.move_to(__lowerCAmelCase			)
         ckpt_arr.append(__lowerCAmelCase			)

         lowercase__    :   Any						= 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(__lowerCAmelCase			)
        self.add(*__lowerCAmelCase     ,       *__lowerCAmelCase			)

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

        lowercase__    :   Optional[Any]						= 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__    :   str						= 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__    :   Union[str, Any]						= 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]			)

        lowercase__    :   Tuple						= [meta_mem.copy() for i in range(6			)]
        lowercase__    :   Any						= [meta_mem.copy() for i in range(6			)]
        lowercase__    :   str						= VGroup(*__lowerCAmelCase			).arrange(__lowerCAmelCase     ,       buff=0			)
        lowercase__    :   Any						= VGroup(*__lowerCAmelCase			).arrange(__lowerCAmelCase     ,       buff=0			)
        lowercase__    :   Union[str, Any]						= VGroup(__lowerCAmelCase     ,       __lowerCAmelCase			).arrange(__lowerCAmelCase     ,       buff=0			)
        lowercase__    :   Union[str, Any]						= Text('''Disk'''     ,       font_size=24			)
        lowercase__    :   Tuple						= Group(__lowerCAmelCase     ,       __lowerCAmelCase			).arrange(__lowerCAmelCase     ,       buff=0.5     ,       aligned_edge=__lowerCAmelCase			)
        disk.move_to([-4.0, -1.2_5, 0]			)
        self.play(Write(__lowerCAmelCase     ,       run_time=3			)     ,       Write(__lowerCAmelCase     ,       run_time=1			)     ,       Create(__lowerCAmelCase     ,       run_time=1			)			)

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

        self.play(FadeOut(__lowerCAmelCase			)			)

        lowercase__    :   Dict						= 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(__lowerCAmelCase     ,       run_time=3			)			)

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

        self.wait()

