

'''simple docstring'''


from manim import *




class     a			(				SCREAMING_SNAKE_CASE       ):






						"""simple docstring"""




						def       __magic_name__							(							self					:      Dict  ):



								'''simple docstring'''



								snake_case__						:       List[str]       						=  Rectangle(height=0.5			,			width=0.5  )
								snake_case__						:       List[str]       						=  Rectangle(height=0.4_6			,			width=0.4_6  ).set_stroke(width=0  )

								snake_case__						:       Tuple       						=  Rectangle(height=0.2_5			,			width=0.2_5  )

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

								snake_case__						:       Dict       						=  [mem.copy() for i in range(4  )]
								snake_case__						:       int       						=  VGroup(*snake_case_  ).arrange(snake_case_			,			buff=0  )
								snake_case__						:       List[str]       						=  Text('''GPU'''			,			font_size=2_4  )
								snake_case__						:       Optional[Any]       						=  Group(snake_case_			,			snake_case_  ).arrange(snake_case_			,			buff=0.5			,			aligned_edge=snake_case_  )
								gpu.move_to([-1, -1, 0]  )
								self.add(snake_case_  )

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

								snake_case__						:       List[str]       						=  Text('''Model'''			,			font_size=2_4  )
								snake_case__						:       Optional[int]       						=  Group(snake_case_			,			snake_case_  ).arrange(snake_case_			,			buff=0.5			,			aligned_edge=snake_case_  )
								model.move_to([3, -1.0, 0]  )
								self.add(snake_case_  )

								snake_case__						:       int       						=  []
								snake_case__						:       List[str]       						=  []

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

										snake_case__						:       Union[str, Any]       						=  Rectangle(height=0.4_6			,			width=0.4_6  ).set_stroke(width=0.0  ).set_fill(snake_case_			,			opacity=0.8  )
										cpu_target.move_to(cpu_left_col_base[i]  )
										model_cpu_arr.append(snake_case_  )

								self.add(*snake_case_			,			*snake_case_  )

								snake_case__						:       Union[str, Any]       						=  [meta_mem.copy() for i in range(6  )]
								snake_case__						:       Any       						=  [meta_mem.copy() for i in range(6  )]
								snake_case__						:       Optional[Any]       						=  VGroup(*snake_case_  ).arrange(snake_case_			,			buff=0  )
								snake_case__						:       Tuple       						=  VGroup(*snake_case_  ).arrange(snake_case_			,			buff=0  )
								snake_case__						:       Union[str, Any]       						=  VGroup(snake_case_			,			snake_case_  ).arrange(snake_case_			,			buff=0  )
								snake_case__						:       Dict       						=  Text('''Disk'''			,			font_size=2_4  )
								snake_case__						:       str       						=  Group(snake_case_			,			snake_case_  ).arrange(snake_case_			,			buff=0.5			,			aligned_edge=snake_case_  )
								disk.move_to([-4, -1.2_5, 0]  )
								self.add(snake_case_			,			snake_case_  )

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

								snake_case__						:       Optional[Any]       						=  MarkupText(
								    F"""<b>Key:</b>\n\n<span fgcolor='{YELLOW}'>●</span> Empty Model"""			,			font_size=1_8			,			)

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

								self.add(snake_case_			,			snake_case_  )

								snake_case__						:       Optional[int]       						=  MarkupText(
								    F"""<span fgcolor='{BLUE}'>●</span> Checkpoint"""			,			font_size=1_8			,			)

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

								snake_case__						:       List[str]       						=  MarkupText(
								    F"""Now watch as an input is passed through the model\nand how the memory is utilized and handled."""			,			font_size=2_4			,			)
								step_a.move_to([2, 2, 0]  )

								self.play(Write(snake_case_  )  )

								snake_case__						:       Any       						=  Square(0.3  )
								input.set_fill(snake_case_			,			opacity=1.0  )
								input.set_stroke(width=0.0  )
								input.next_to(model_base[0]			,			snake_case_			,			buff=0.5  )

								self.play(Write(snake_case_  )  )

								input.generate_target()
								input.target.next_to(model_arr[0]			,			direction=snake_case_			,			buff=0.0_2  )
								self.play(MoveToTarget(snake_case_  )  )

								self.play(FadeOut(snake_case_  )  )

								snake_case__						:       Union[str, Any]       						=  Arrow(start=snake_case_			,			end=snake_case_			,			color=snake_case_			,			buff=0.5  )
								a.next_to(model_arr[0].get_left()			,			snake_case_			,			buff=0.2  )

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

								snake_case__						:       Union[str, 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=2_4			,			)
								step_a.move_to([2, 2, 0]  )

								self.play(Write(snake_case_			,			run_time=3  )  )

								snake_case__						:       Optional[Any]       						=  {'''run_time''': 1, '''fade_in''': True, '''fade_out''': True, '''buff''': 0.0_2}

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

								snake_case__						:       Tuple       						=  a.copy()
								for i in range(6  ):
										a_c.next_to(model_arr[i].get_right() + 0.0_2			,			snake_case_			,			buff=0.2  )

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

										snake_case__						:       List[str]       						=  AnimationGroup(
										    FadeOut(snake_case_			,			run_time=0.5  )			,			MoveToTarget(snake_case_			,			run_time=0.5  )			,			FadeIn(snake_case_			,			run_time=0.5  )			,			lag_ratio=0.2  )

										self.play(snake_case_  )

										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:
														snake_case__						:       str       						=  0.7

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

												self.play(
												    Circumscribe(model_arr[-1]			,			color=snake_case_			,			**snake_case_  )			,			Circumscribe(cpu_left_col_base[-1]			,			color=snake_case_			,			**snake_case_  )			,			Circumscribe(gpu_rect[0]			,			color=snake_case_			,			**snake_case_  )			,			)

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

										snake_case__						:       List[Any]       						=  a_c
										snake_case__						:       Union[str, Any]       						=  a_c.copy()

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

								snake_case__						:       List[str]       						=  MarkupText(F"""Inference on a model too large for GPU memory\nis successfully completed."""			,			font_size=2_4  )
								step_a.move_to([2, 2, 0]  )

								self.play(Write(snake_case_			,			run_time=3  )			,			MoveToTarget(snake_case_  )  )

								self.wait()






'''simple docstring'''


import argparse

from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection

from diffusers import UnCLIPImageVariationPipeline, UnCLIPPipeline


if __name__ == "__main__":
					lowerCAmelCase__						:       Dict										=						argparse.ArgumentParser()

					parser.add_argument("""--dump_path""", default=None, type=str, required=True, help="""Path to the output model.""")

					parser.add_argument(
					    """--txt2img_unclip""",
					    default="""kakaobrain/karlo-v1-alpha""",
					    type=str,
					    required=False,
					    help="""The pretrained txt2img unclip.""",
					)

					lowerCAmelCase__						:       List[Any]										=						parser.parse_args()

					lowerCAmelCase__						:       Optional[Any]										=						UnCLIPPipeline.from_pretrained(args.txtaimg_unclip)

					lowerCAmelCase__						:       Optional[int]										=						CLIPImageProcessor()
					lowerCAmelCase__						:       str										=						CLIPVisionModelWithProjection.from_pretrained("""openai/clip-vit-large-patch14""")

					lowerCAmelCase__						:       Dict										=						UnCLIPImageVariationPipeline(
					    decoder=txtaimg.decoder,
					    text_encoder=txtaimg.text_encoder,
					    tokenizer=txtaimg.tokenizer,
					    text_proj=txtaimg.text_proj,
					    feature_extractor=feature_extractor,
					    image_encoder=image_encoder,
					    super_res_first=txtaimg.super_res_first,
					    super_res_last=txtaimg.super_res_last,
					    decoder_scheduler=txtaimg.decoder_scheduler,
					    super_res_scheduler=txtaimg.super_res_scheduler,
					)

					imgaimg.save_pretrained(args.dump_path)



