import gc
import random
import unittest

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

from diffusers import (
    DiffusionPipeline,
    UnCLIPImageVariationPipeline,
    UnCLIPScheduler,
    UNetaDConditionModel,
    UNetaDModel,
)
from diffusers.pipelines.unclip.text_proj import UnCLIPTextProjModel
from diffusers.utils import floats_tensor, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, load_image, require_torch_gpu, skip_mps

from ..pipeline_params import IMAGE_VARIATION_BATCH_PARAMS, IMAGE_VARIATION_PARAMS
from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()






class      UpperCamelCase    (						__snake_case       ,    unittest.TestCase     ):
	lowerCAmelCase	:   Optional[int]				=	UnCLIPImageVariationPipeline
	lowerCAmelCase	:   Any				=	IMAGE_VARIATION_PARAMS - {'height', 'width', 'guidance_scale'}
	lowerCAmelCase	:   str				=	IMAGE_VARIATION_BATCH_PARAMS

	lowerCAmelCase	:   Optional[int]				=	[
	    'generator',
	    'return_dict',
	    'decoder_num_inference_steps',
	    'super_res_num_inference_steps',
	]
	lowerCAmelCase	:   Union[str, 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__							=       CLIPVisionConfig(
							    hidden_size=self.text_embedder_hidden_size			,	projection_dim=self.text_embedder_hidden_size			,	num_hidden_layers=5			,	num_attention_heads=4			,	image_size=32			,	intermediate_size=37			,	patch_size=1			,	)
							return CLIPVisionModelWithProjection(UpperCAmelCase__		)



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

							A__							=       {
							    "clip_embeddings_dim": self.text_embedder_hidden_size,
							    "time_embed_dim": self.time_embed_dim,
							    "cross_attention_dim": self.cross_attention_dim,
							}

							A__							=       UnCLIPTextProjModel(**UpperCAmelCase__		)
							return model



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

							A__							=       {
							    "sample_size": 32,
							    # RGB in channels
							    "in_channels": 3,
							    # Out channels is double in channels because predicts mean and variance
							    "out_channels": 6,
							    "down_block_types": ("ResnetDownsampleBlock2D", "SimpleCrossAttnDownBlock2D"),
							    "up_block_types": ("SimpleCrossAttnUpBlock2D", "ResnetUpsampleBlock2D"),
							    "mid_block_type": "UNetMidBlock2DSimpleCrossAttn",
							    "block_out_channels": (self.block_out_channels_a, self.block_out_channels_a * 2),
							    "layers_per_block": 1,
							    "cross_attention_dim": self.cross_attention_dim,
							    "attention_head_dim": 4,
							    "resnet_time_scale_shift": "scale_shift",
							    "class_embed_type": "identity",
							}

							A__							=       UNetaDConditionModel(**UpperCAmelCase__		)
							return model



	@property
	def 			__A							(       self		):
							return {
							    "sample_size": 64,
							    "layers_per_block": 1,
							    "down_block_types": ("ResnetDownsampleBlock2D", "ResnetDownsampleBlock2D"),
							    "up_block_types": ("ResnetUpsampleBlock2D", "ResnetUpsampleBlock2D"),
							    "block_out_channels": (self.block_out_channels_a, self.block_out_channels_a * 2),
							    "in_channels": 6,
							    "out_channels": 3,
							}



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

							A__							=       UNetaDModel(**self.dummy_super_res_kwargs		)
							return model



	@property
	def 			__A							(       self		):
							# seeded differently to get different unet than `self.dummy_super_res_first`
							torch.manual_seed(1		)

							A__							=       UNetaDModel(**self.dummy_super_res_kwargs		)
							return model



	def 			__A							(       self		):
							A__							=       self.dummy_decoder
							A__							=       self.dummy_text_proj
							A__							=       self.dummy_text_encoder
							A__							=       self.dummy_tokenizer
							A__							=       self.dummy_super_res_first
							A__							=       self.dummy_super_res_last

							A__							=       UnCLIPScheduler(
							    variance_type="learned_range"			,	prediction_type="epsilon"			,	num_train_timesteps=1_000			,	)

							A__							=       UnCLIPScheduler(
							    variance_type="fixed_small_log"			,	prediction_type="epsilon"			,	num_train_timesteps=1_000			,	)

							A__							=       CLIPImageProcessor(crop_size=32			,	size=32		)

							A__							=       self.dummy_image_encoder

							return {
							    "decoder": decoder,
							    "text_encoder": text_encoder,
							    "tokenizer": tokenizer,
							    "text_proj": text_proj,
							    "feature_extractor": feature_extractor,
							    "image_encoder": image_encoder,
							    "super_res_first": super_res_first,
							    "super_res_last": super_res_last,
							    "decoder_scheduler": decoder_scheduler,
							    "super_res_scheduler": super_res_scheduler,
							}



	def 			__A							(       self			,	UpperCAmelCase__			,	UpperCAmelCase__=0			,	UpperCAmelCase__=True		):
							A__							=       floats_tensor((1, 3, 32, 32)			,	rng=random.Random(UpperCAmelCase__		)		).to(UpperCAmelCase__		)
							if str(UpperCAmelCase__		).startswith("mps"		):
													A__							=       torch.manual_seed(UpperCAmelCase__		)
							else:
													A__							=       torch.Generator(device=UpperCAmelCase__		).manual_seed(UpperCAmelCase__		)

							if pil_image:
													A__							=       input_image * 0.5 + 0.5
													A__							=       input_image.clamp(0			,	1		)
													A__							=       input_image.cpu().permute(0			,	2			,	3			,	1		).float().numpy()
													A__							=       DiffusionPipeline.numpy_to_pil(UpperCAmelCase__		)[0]

							return {
							    "image": input_image,
							    "generator": generator,
							    "decoder_num_inference_steps": 2,
							    "super_res_num_inference_steps": 2,
							    "output_type": "np",
							}



	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__							=       self.get_dummy_inputs(UpperCAmelCase__			,	pil_image=UpperCAmelCase__		)

							A__							=       pipe(**UpperCAmelCase__		)
							A__							=       output.images

							A__							=       self.get_dummy_inputs(UpperCAmelCase__			,	pil_image=UpperCAmelCase__		)

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

							A__							=       image[0, -3:, -3:, -1]
							A__							=       image_from_tuple[0, -3:, -3:, -1]

							assert image.shape == (1, 64, 64, 3)

							A__							=       np.array(
							    [
							        0.9_997,
							        0.0_002,
							        0.9_997,
							        0.9_997,
							        0.9_969,
							        0.0_023,
							        0.9_997,
							        0.9_969,
							        0.9_970,
							    ]		)

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



	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__							=       self.get_dummy_inputs(UpperCAmelCase__			,	pil_image=UpperCAmelCase__		)

							A__							=       pipe(**UpperCAmelCase__		)
							A__							=       output.images

							A__							=       self.get_dummy_inputs(UpperCAmelCase__			,	pil_image=UpperCAmelCase__		)

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

							A__							=       image[0, -3:, -3:, -1]
							A__							=       image_from_tuple[0, -3:, -3:, -1]

							assert image.shape == (1, 64, 64, 3)

							A__							=       np.array([0.9_997, 0.0_003, 0.9_997, 0.9_997, 0.9_970, 0.0_024, 0.9_997, 0.9_971, 0.9_971]		)

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



	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__							=       self.get_dummy_inputs(UpperCAmelCase__			,	pil_image=UpperCAmelCase__		)
							A__							=       [
							    pipeline_inputs["image"],
							    pipeline_inputs["image"],
							]

							A__							=       pipe(**UpperCAmelCase__		)
							A__							=       output.images

							A__							=       self.get_dummy_inputs(UpperCAmelCase__			,	pil_image=UpperCAmelCase__		)
							A__							=       [
							    tuple_pipeline_inputs["image"],
							    tuple_pipeline_inputs["image"],
							]

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

							A__							=       image[0, -3:, -3:, -1]
							A__							=       image_from_tuple[0, -3:, -3:, -1]

							assert image.shape == (2, 64, 64, 3)

							A__							=       np.array(
							    [
							        0.9_997,
							        0.9_989,
							        0.0_008,
							        0.0_021,
							        0.9_960,
							        0.0_018,
							        0.0_014,
							        0.0_002,
							        0.9_933,
							    ]		)

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



	def 			__A							(       self		):
							A__							=       torch.device("cpu"		)






							class      UpperCamelCase    :
								lowerCAmelCase	:   Any				=	1





							A__							=       self.get_dummy_components()

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

							pipe.set_progress_bar_config(disable=UpperCAmelCase__		)

							A__							=       torch.Generator(device=UpperCAmelCase__		).manual_seed(0		)
							A__							=       pipe.decoder.dtype
							A__							=       1

							A__							=       (
							    batch_size,
							    pipe.decoder.config.in_channels,
							    pipe.decoder.config.sample_size,
							    pipe.decoder.config.sample_size,
							)
							A__							=       pipe.prepare_latents(
							    UpperCAmelCase__			,	dtype=UpperCAmelCase__			,	device=UpperCAmelCase__			,	generator=UpperCAmelCase__			,	latents=UpperCAmelCase__			,	scheduler=DummyScheduler()		)

							A__							=       (
							    batch_size,
							    pipe.super_res_first.config.in_channels // 2,
							    pipe.super_res_first.config.sample_size,
							    pipe.super_res_first.config.sample_size,
							)
							A__							=       pipe.prepare_latents(
							    UpperCAmelCase__			,	dtype=UpperCAmelCase__			,	device=UpperCAmelCase__			,	generator=UpperCAmelCase__			,	latents=UpperCAmelCase__			,	scheduler=DummyScheduler()		)

							A__							=       self.get_dummy_inputs(UpperCAmelCase__			,	pil_image=UpperCAmelCase__		)

							A__							=       pipe(
							    **UpperCAmelCase__			,	decoder_latents=UpperCAmelCase__			,	super_res_latents=UpperCAmelCase__		).images

							A__							=       self.get_dummy_inputs(UpperCAmelCase__			,	pil_image=UpperCAmelCase__		)
							# Don't pass image, instead pass embedding
							A__							=       pipeline_inputs.pop("image"		)
							A__							=       pipe.image_encoder(UpperCAmelCase__		).image_embeds

							A__							=       pipe(
							    **UpperCAmelCase__			,	decoder_latents=UpperCAmelCase__			,	super_res_latents=UpperCAmelCase__			,	image_embeddings=UpperCAmelCase__			,	).images

							# make sure passing text embeddings manually is identical
							assert np.abs(img_out_a - img_out_a		).max() < 1e-4



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

							# Check is relaxed because there is not a torch 2.0 sliced attention added kv processor
							A__							=       1e-2

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



	@skip_mps
	def 			__A							(       self		):
							A__							=       torch_device == "cpu"
							A__							=       True
							A__							=       [
							    "decoder_num_inference_steps",
							    "super_res_num_inference_steps",
							]

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



	def 			__A							(       self		):
							A__							=       [
							    "decoder_num_inference_steps",
							    "super_res_num_inference_steps",
							]

							if torch_device == "mps":
													# TODO: MPS errors with larger batch sizes
													A__							=       [2, 3]
													self._test_inference_batch_consistent(
													    batch_sizes=UpperCAmelCase__			,	additional_params_copy_to_batched_inputs=UpperCAmelCase__			,	)
							else:
													self._test_inference_batch_consistent(
													    additional_params_copy_to_batched_inputs=UpperCAmelCase__		)



	@skip_mps
	def 			__A							(       self		):
							return super().test_dict_tuple_outputs_equivalent()



	@skip_mps
	def 			__A							(       self		):
							return super().test_save_load_local()



	@skip_mps
	def 			__A							(       self		):
							return super().test_save_load_optional_components()






@slow
@require_torch_gpu
class      UpperCamelCase    (						unittest.TestCase     ):



	def 			__A							(       self		):
							# clean up the VRAM after each test
							super().tearDown()
							gc.collect()
							torch.cuda.empty_cache()



	def 			__A							(       self		):
							A__							=       load_image(
							    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/unclip/cat.png"		)
							A__							=       load_numpy(
							    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
							    "/unclip/karlo_v1_alpha_cat_variation_fp16.npy"		)

							A__							=       UnCLIPImageVariationPipeline.from_pretrained(
							    "kakaobrain/karlo-v1-alpha-image-variations"			,	torch_dtype=torch.floataa		)
							A__							=       pipeline.to(UpperCAmelCase__		)
							pipeline.set_progress_bar_config(disable=UpperCAmelCase__		)

							A__							=       torch.Generator(device="cpu"		).manual_seed(0		)
							A__							=       pipeline(
							    UpperCAmelCase__			,	generator=UpperCAmelCase__			,	output_type="np"			,	)

							A__							=       output.images[0]

							assert image.shape == (256, 256, 3)

							assert_mean_pixel_difference(UpperCAmelCase__			,	UpperCAmelCase__			,	15		)










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





