



import unittest

import numpy as np
import torch
from torch import nn
from transformers import (
    CLIPImageProcessor,
    CLIPTextConfig,
    CLIPTextModelWithProjection,
    CLIPTokenizer,
    CLIPVisionConfig,
    CLIPVisionModelWithProjection,
)

from diffusers import KandinskyVaaPriorPipeline, PriorTransformer, UnCLIPScheduler
from diffusers.utils import torch_device
from diffusers.utils.testing_utils import enable_full_determinism, skip_mps

from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()






class      UpperCamelCase    (						_UpperCAmelCase       ,    unittest.TestCase     ):
	lowerCAmelCase	:   str				=	KandinskyVaaPriorPipeline
	lowerCAmelCase	:   int				=	["""prompt"""]
	lowerCAmelCase	:   str				=	["""prompt""", """negative_prompt"""]
	lowerCAmelCase	:   str				=	[
	    """num_images_per_prompt""",
	    """generator""",
	    """num_inference_steps""",
	    """latents""",
	    """negative_prompt""",
	    """guidance_scale""",
	    """output_type""",
	    """return_dict""",
	]
	lowerCAmelCase	:   List[Any]				=	False



	@property
	def 			__A							(       self		):
							return 32



	@property
	def 			__A							(       self		):
							return 32



	@property
	def 			__A							(       self		):
							return self.time_input_dim



	@property
	def 			__A							(       self		):
							return self.time_input_dim * 4



	@property
	def 			__A							(       self		):
							return 100



	@property
	def 			__A							(       self		):
							A__							=       CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"		)
							return tokenizer



	@property
	def 			__A							(       self		):
							torch.manual_seed(0		)
							A__							=       CLIPTextConfig(
							    bos_token_id=0			,	eos_token_id=2			,	hidden_size=self.text_embedder_hidden_size			,	projection_dim=self.text_embedder_hidden_size			,	intermediate_size=37			,	layer_norm_eps=1e-0_5			,	num_attention_heads=4			,	num_hidden_layers=5			,	pad_token_id=1			,	vocab_size=1_000			,	)
							return CLIPTextModelWithProjection(UpperCAmelCase__		)



	@property
	def 			__A							(       self		):
							torch.manual_seed(0		)

							A__							=       {
							    "num_attention_heads": 2,
							    "attention_head_dim": 12,
							    "embedding_dim": self.text_embedder_hidden_size,
							    "num_layers": 1,
							}

							A__							=       PriorTransformer(**UpperCAmelCase__		)
							# clip_std and clip_mean is initialized to be 0 so PriorTransformer.post_process_latents will always return 0 - set clip_std to be 1 so it won't return 0
							A__							=       nn.Parameter(torch.ones(model.clip_std.shape		)		)
							return model



	@property
	def 			__A							(       self		):
							torch.manual_seed(0		)
							A__							=       CLIPVisionConfig(
							    hidden_size=self.text_embedder_hidden_size			,	image_size=224			,	projection_dim=self.text_embedder_hidden_size			,	intermediate_size=37			,	num_attention_heads=4			,	num_channels=3			,	num_hidden_layers=5			,	patch_size=14			,	)

							A__							=       CLIPVisionModelWithProjection(UpperCAmelCase__		)
							return model



	@property
	def 			__A							(       self		):
							A__							=       CLIPImageProcessor(
							    crop_size=224			,	do_center_crop=UpperCAmelCase__			,	do_normalize=UpperCAmelCase__			,	do_resize=UpperCAmelCase__			,	image_mean=[0.48_145_466, 0.4_578_275, 0.40_821_073]			,	image_std=[0.26_862_954, 0.26_130_258, 0.27_577_711]			,	resample=3			,	size=224			,	)

							return image_processor



	def 			__A							(       self		):
							A__							=       self.dummy_prior
							A__							=       self.dummy_image_encoder
							A__							=       self.dummy_text_encoder
							A__							=       self.dummy_tokenizer
							A__							=       self.dummy_image_processor

							A__							=       UnCLIPScheduler(
							    variance_type="fixed_small_log"			,	prediction_type="sample"			,	num_train_timesteps=1_000			,	clip_sample=UpperCAmelCase__			,	clip_sample_range=10.0			,	)

							A__							=       {
							    "prior": prior,
							    "image_encoder": image_encoder,
							    "text_encoder": text_encoder,
							    "tokenizer": tokenizer,
							    "scheduler": scheduler,
							    "image_processor": image_processor,
							}

							return components



	def 			__A							(       self			,	UpperCAmelCase__			,	UpperCAmelCase__=0		):
							if str(UpperCAmelCase__		).startswith("mps"		):
													A__							=       torch.manual_seed(UpperCAmelCase__		)
							else:
													A__							=       torch.Generator(device=UpperCAmelCase__		).manual_seed(UpperCAmelCase__		)
							A__							=       {
							    "prompt": "horse",
							    "generator": generator,
							    "guidance_scale": 4.0,
							    "num_inference_steps": 2,
							    "output_type": "np",
							}
							return inputs



	def 			__A							(       self		):
							A__							=       "cpu"

							A__							=       self.get_dummy_components()

							A__							=       self.pipeline_class(**UpperCAmelCase__		)
							A__							=       pipe.to(UpperCAmelCase__		)

							pipe.set_progress_bar_config(disable=UpperCAmelCase__		)

							A__							=       pipe(**self.get_dummy_inputs(UpperCAmelCase__		)		)
							A__							=       output.image_embeds

							A__							=       pipe(
							    **self.get_dummy_inputs(UpperCAmelCase__		)			,	return_dict=UpperCAmelCase__			,	)[0]

							A__							=       image[0, -10:]
							A__							=       image_from_tuple[0, -10:]

							assert image.shape == (1, 32)

							A__							=       np.array(
							    [-0.0_532, 1.7_120, 0.3_656, -1.0_852, -0.8_946, -1.1_756, 0.4_348, 0.2_482, 0.5_146, -0.1_156]		)

							assert np.abs(image_slice.flatten() - expected_slice		).max() < 1e-2
							assert np.abs(image_from_tuple_slice.flatten() - expected_slice		).max() < 1e-2



	@skip_mps
	def 			__A							(       self		):
							A__							=       torch_device == "cpu"
							A__							=       True
							A__							=       False

							self._test_inference_batch_single_identical(
							    test_max_difference=UpperCAmelCase__			,	relax_max_difference=UpperCAmelCase__			,	test_mean_pixel_difference=UpperCAmelCase__			,	)



	@skip_mps
	def 			__A							(       self		):
							A__							=       torch_device == "cpu"
							A__							=       False

							self._test_attention_slicing_forward_pass(
							    test_max_difference=UpperCAmelCase__			,	test_mean_pixel_difference=UpperCAmelCase__			,	)










from manim import *






class      UpperCamelCase    (						_UpperCAmelCase     ):



	def 			__A							(       self		):
							A__							=       Rectangle(height=0.5			,	width=0.5		)
							A__							=       Rectangle(height=0.46			,	width=0.46		).set_stroke(width=0		)

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

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

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

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

							A__							=       []
							for i, rect in enumerate(UpperCAmelCase__		):
													rect.set_stroke(UpperCAmelCase__		)
													# target = fill.copy().set_fill(YELLOW, opacity=0.7)
													# target.move_to(rect)
													# self.add(target)

													A__							=       Rectangle(height=0.46 / 4			,	width=0.46 / 3		).set_stroke(width=0.0		).set_fill(UpperCAmelCase__			,	opacity=0.7		)

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

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

							A__							=       Text("Loaded Checkpoint"			,	font_size=24		)
							A__							=       Group(UpperCAmelCase__			,	UpperCAmelCase__		).arrange(UpperCAmelCase__			,	aligned_edge=UpperCAmelCase__			,	buff=0.4		)
							checkpoint.move_to([3, 0.5, 0]		)

							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(UpperCAmelCase__			,	UpperCAmelCase__		)

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

							blue_text.next_to(UpperCAmelCase__			,	DOWN * 2.4			,	aligned_edge=key_text.get_left()		)

							A__							=       MarkupText(
							    F"""Next, a <i><span fgcolor=\"{BLUE}\">second</span></i> model is loaded into memory,\nwith the weights of a <span fgcolor=\"{BLUE}\">single shard</span>."""			,	font_size=24			,	)
							step_a.move_to([2, 2, 0]		)
							self.play(Write(UpperCAmelCase__		)			,	Write(UpperCAmelCase__		)		)

							self.play(Write(UpperCAmelCase__			,	run_time=1		)			,	Create(UpperCAmelCase__			,	run_time=1		)		)

							A__							=       []
							A__							=       []
							for i, rect in enumerate(UpperCAmelCase__		):
													A__							=       fill.copy().set_fill(UpperCAmelCase__			,	opacity=0.7		)
													target.move_to(UpperCAmelCase__		)
													first_animations.append(GrowFromCenter(UpperCAmelCase__			,	run_time=1		)		)

													A__							=       target.copy()
													cpu_target.generate_target()
													if i < 5:
																			cpu_target.target.move_to(cpu_left_col_base[i + 1]		)
													else:
																			cpu_target.target.move_to(cpu_right_col_base[i - 5]		)
													second_animations.append(MoveToTarget(UpperCAmelCase__			,	run_time=1.5		)		)

							self.play(*UpperCAmelCase__		)
							self.play(*UpperCAmelCase__		)
							self.wait()





