
"""simple docstring"""




import torch
from torch import nn
from transformers import CLIPPreTrainedModel, CLIPVisionModel

from ...models.attention import BasicTransformerBlock
from ...utils import logging


UpperCAmelCase__   :					Any												=				logging.get_logger(__name__)  # pylint: disable=invalid-name



class 			lowerCAmelCase_		(a__						):


				"""simple docstring"""



				def __init__(self     , SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__=7_68					)      ->	int:



											"""simple docstring"""




											super().__init__(SCREAMING_SNAKE_CASE__					)
											SCREAMING_SNAKE_CASE__		: Any     =     proj_size

											SCREAMING_SNAKE_CASE__		: Tuple     =     CLIPVisionModel(SCREAMING_SNAKE_CASE__					)
											SCREAMING_SNAKE_CASE__		: List[str]     =     PaintByExampleMapper(SCREAMING_SNAKE_CASE__					)
											SCREAMING_SNAKE_CASE__		: str     =     nn.LayerNorm(config.hidden_size					)
											SCREAMING_SNAKE_CASE__		: List[Any]     =     nn.Linear(config.hidden_size     , self.proj_size					)

											# uncondition for scaling
											SCREAMING_SNAKE_CASE__		: List[Any]     =     nn.Parameter(torch.randn((1, 1, self.proj_size)					)					)



				def 			__magic_name__					(self     , SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__=False					)      ->	str:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__		: List[Any]     =     self.model(pixel_values=SCREAMING_SNAKE_CASE__					)
											SCREAMING_SNAKE_CASE__		: Dict     =     clip_output.pooler_output
											SCREAMING_SNAKE_CASE__		: Union[str, Any]     =     self.mapper(latent_states[:, None]					)
											SCREAMING_SNAKE_CASE__		: Union[str, Any]     =     self.final_layer_norm(SCREAMING_SNAKE_CASE__					)
											SCREAMING_SNAKE_CASE__		: str     =     self.proj_out(SCREAMING_SNAKE_CASE__					)
											if return_uncond_vector:
																		return latent_states, self.uncond_vector

											return latent_states






class 			lowerCAmelCase_		(nn.Module						):


				"""simple docstring"""



				def __init__(self     , SCREAMING_SNAKE_CASE__					)      ->	Tuple:



											"""simple docstring"""




											super().__init__()
											SCREAMING_SNAKE_CASE__		: int     =     (config.num_hidden_layers + 1) // 5
											SCREAMING_SNAKE_CASE__		: int     =     config.hidden_size
											SCREAMING_SNAKE_CASE__		: Optional[Any]     =     1
											SCREAMING_SNAKE_CASE__		: Any     =     nn.ModuleList(
											    [
											        BasicTransformerBlock(SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__     , activation_fn="""gelu"""     , attention_bias=SCREAMING_SNAKE_CASE__					)
											        for _ in range(SCREAMING_SNAKE_CASE__					)
											    ]					)



				def 			__magic_name__					(self     , SCREAMING_SNAKE_CASE__					)      ->	Any:



											"""simple docstring"""




											for block in self.blocks:
																		SCREAMING_SNAKE_CASE__		: Tuple     =     block(SCREAMING_SNAKE_CASE__					)

											return hidden_states






"""simple docstring"""




from manim import *



class 			lowerCAmelCase_		(a__						):


				"""simple docstring"""



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



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__		: Dict     =     Rectangle(height=0.5     , width=0.5					)
											SCREAMING_SNAKE_CASE__		: Union[str, Any]     =     Rectangle(height=0.46     , width=0.46					).set_stroke(width=0					)

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

											SCREAMING_SNAKE_CASE__		: Optional[int]     =     [mem.copy() for i in range(1					)]
											SCREAMING_SNAKE_CASE__		: Optional[Any]     =     VGroup(*SCREAMING_SNAKE_CASE__					).arrange(SCREAMING_SNAKE_CASE__     , buff=0					)
											SCREAMING_SNAKE_CASE__		: str     =     Text("""GPU"""     , font_size=24					)
											SCREAMING_SNAKE_CASE__		: str     =     Group(SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__					).arrange(SCREAMING_SNAKE_CASE__     , buff=0.5     , aligned_edge=SCREAMING_SNAKE_CASE__					)
											gpu.align_to(SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__					)
											gpu.set_x(gpu.get_x() - 1					)

											self.add(SCREAMING_SNAKE_CASE__					)

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

											SCREAMING_SNAKE_CASE__		: Optional[Any]     =     Text("""Model"""     , font_size=24					)
											SCREAMING_SNAKE_CASE__		: Any     =     Group(SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__					).arrange(SCREAMING_SNAKE_CASE__     , buff=0.5     , aligned_edge=SCREAMING_SNAKE_CASE__					)
											model.move_to([3, -1.0, 0]					)

											self.play(
											    Create(SCREAMING_SNAKE_CASE__     , run_time=1					)     , Create(SCREAMING_SNAKE_CASE__     , run_time=1					)     , Create(SCREAMING_SNAKE_CASE__     , run_time=1					)     , )

											SCREAMING_SNAKE_CASE__		: Tuple     =     MarkupText(
											    F'''First, an empty model skeleton is loaded\ninto <span fgcolor=\'{YELLOW}\'>memory</span> without using much RAM.'''     , font_size=24     , )

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

											SCREAMING_SNAKE_CASE__		: Tuple     =     MarkupText(
											    F'''<b>Key:</b>\n\n<span fgcolor=\'{YELLOW}\'>●</span> Empty Model'''     , font_size=18     , )

											key_text.move_to([-5, 2.4, 0]					)

											step_a.move_to([2, 2, 0]					)
											self.play(Write(SCREAMING_SNAKE_CASE__     , run_time=2.5					)     , Write(SCREAMING_SNAKE_CASE__					)     , Write(SCREAMING_SNAKE_CASE__					)					)

											self.add(SCREAMING_SNAKE_CASE__					)

											SCREAMING_SNAKE_CASE__		: Optional[int]     =     []
											SCREAMING_SNAKE_CASE__		: Union[str, Any]     =     []
											SCREAMING_SNAKE_CASE__		: str     =     []
											for i, rect in enumerate(SCREAMING_SNAKE_CASE__					):
																		SCREAMING_SNAKE_CASE__		: Union[str, Any]     =     Rectangle(height=0.46     , width=0.46					).set_stroke(width=0.0					).set_fill(SCREAMING_SNAKE_CASE__     , opacity=0.7					)
																		cpu_target.move_to(SCREAMING_SNAKE_CASE__					)
																		cpu_target.generate_target()
																		SCREAMING_SNAKE_CASE__		: Tuple     =     0.46 / 4
																		SCREAMING_SNAKE_CASE__		: int     =     0.46 / 3

																		if i == 0:
																									cpu_target.target.next_to(cpu_left_col_base[0].get_corner(DOWN + LEFT					)     , buff=0.02     , direction=SCREAMING_SNAKE_CASE__					)
																									cpu_target.target.set_x(cpu_target.target.get_x() + 0.1					)
																		elif i == 3:
																									cpu_target.target.next_to(cpu_targs[0].target     , direction=SCREAMING_SNAKE_CASE__     , buff=0.0					)
																		else:
																									cpu_target.target.next_to(cpu_targs[i - 1].target     , direction=SCREAMING_SNAKE_CASE__     , buff=0.0					)
																		cpu_targs.append(SCREAMING_SNAKE_CASE__					)

																		first_animations.append(rect.animate(run_time=0.5					).set_stroke(SCREAMING_SNAKE_CASE__					)					)
																		second_animations.append(MoveToTarget(SCREAMING_SNAKE_CASE__     , run_time=1.5					)					)

											self.play(*SCREAMING_SNAKE_CASE__					)
											self.play(*SCREAMING_SNAKE_CASE__					)

											self.wait()




