


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 							A   (	UpperCAmelCase__			,  unittest.TestCase							):
					'''simple docstring'''






					A__      =			UnCLIPImageVariationPipeline
					A__      =			IMAGE_VARIATION_PARAMS - {'''height''', '''width''', '''guidance_scale'''}
					A__      =			IMAGE_VARIATION_BATCH_PARAMS

					A__      =			[
					    '''generator''',
					    '''return_dict''',
					    '''decoder_num_inference_steps''',
					    '''super_res_num_inference_steps''',
					]
					A__      =			False

					@property
					def   lowerCamelCase__					(self   :						Optional[Any]		)						->       Optional[Any]:




							"""simple docstring"""


							return 32

					@property
					def   lowerCamelCase__					(self   :						Optional[Any]		)						->       Optional[int]:




							"""simple docstring"""


							return 32

					@property
					def   lowerCamelCase__					(self   :						Union[str, Any]		)						->       List[Any]:




							"""simple docstring"""


							return self.time_input_dim

					@property
					def   lowerCamelCase__					(self   :						List[Any]		)						->       Optional[int]:




							"""simple docstring"""


							return self.time_input_dim * 4

					@property
					def   lowerCamelCase__					(self   :						Optional[Any]		)						->       str:




							"""simple docstring"""


							return 100

					@property
					def   lowerCamelCase__					(self   :						List[Any]		)						->       str:




							"""simple docstring"""


							lowercase__					      =		CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""		)
							return tokenizer

					@property
					def   lowerCamelCase__					(self   :						List[Any]		)						->       str:




							"""simple docstring"""


							torch.manual_seed(0		)
							lowercase__					      =		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=1000				, )
							return CLIPTextModelWithProjection(_UpperCAmelCase		)

					@property
					def   lowerCamelCase__					(self   :						Tuple		)						->       List[Any]:




							"""simple docstring"""


							torch.manual_seed(0		)
							lowercase__					      =		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   lowerCamelCase__					(self   :						Any		)						->       Tuple:




							"""simple docstring"""


							torch.manual_seed(0		)

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

							lowercase__					      =		UnCLIPTextProjModel(**_UpperCAmelCase		)
							return model

					@property
					def   lowerCamelCase__					(self   :						Any		)						->       List[str]:




							"""simple docstring"""


							torch.manual_seed(0		)

							lowercase__					      =		{
							    """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""",
							}

							lowercase__					      =		UNetaDConditionModel(**_UpperCAmelCase		)
							return model

					@property
					def   lowerCamelCase__					(self   :						Any		)						->       Optional[int]:




							"""simple docstring"""


							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   lowerCamelCase__					(self   :						List[str]		)						->       str:




							"""simple docstring"""


							torch.manual_seed(0		)

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

					@property
					def   lowerCamelCase__					(self   :						Dict		)						->       Union[str, Any]:




							"""simple docstring"""


							torch.manual_seed(1		)

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

					def   lowerCamelCase__					(self   :						Union[str, Any]		)						->       Any:




							"""simple docstring"""


							lowercase__					      =		self.dummy_decoder
							lowercase__					      =		self.dummy_text_proj
							lowercase__					      =		self.dummy_text_encoder
							lowercase__					      =		self.dummy_tokenizer
							lowercase__					      =		self.dummy_super_res_first
							lowercase__					      =		self.dummy_super_res_last

							lowercase__					      =		UnCLIPScheduler(
							    variance_type="""learned_range"""				, prediction_type="""epsilon"""				, num_train_timesteps=1000				, )

							lowercase__					      =		UnCLIPScheduler(
							    variance_type="""fixed_small_log"""				, prediction_type="""epsilon"""				, num_train_timesteps=1000				, )

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

							lowercase__					      =		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   lowerCamelCase__					(self   :						Optional[int]				, _UpperCAmelCase   :						List[str]				, _UpperCAmelCase   :						List[str]=0				, _UpperCAmelCase   :						str=True		)						->       Optional[Any]:




							"""simple docstring"""


							lowercase__					      =		floats_tensor((1, 3, 32, 32)				, rng=random.Random(_UpperCAmelCase		)		).to(_UpperCAmelCase		)
							if str(_UpperCAmelCase		).startswith("""mps"""		):
									lowercase__					      =		torch.manual_seed(_UpperCAmelCase		)
							else:
									lowercase__					      =		torch.Generator(device=_UpperCAmelCase		).manual_seed(_UpperCAmelCase		)

							if pil_image:
									lowercase__					      =		input_image * 0.5 + 0.5
									lowercase__					      =		input_image.clamp(0				, 1		)
									lowercase__					      =		input_image.cpu().permute(0				, 2				, 3				, 1		).float().numpy()
									lowercase__					      =		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   lowerCamelCase__					(self   :						Union[str, Any]		)						->       str:




							"""simple docstring"""


							lowercase__					      =		"""cpu"""

							lowercase__					      =		self.get_dummy_components()

							lowercase__					      =		self.pipeline_class(**_UpperCAmelCase		)
							lowercase__					      =		pipe.to(_UpperCAmelCase		)

							pipe.set_progress_bar_config(disable=_UpperCAmelCase		)

							lowercase__					      =		self.get_dummy_inputs(_UpperCAmelCase				, pil_image=_UpperCAmelCase		)

							lowercase__					      =		pipe(**_UpperCAmelCase		)
							lowercase__					      =		output.images

							lowercase__					      =		self.get_dummy_inputs(_UpperCAmelCase				, pil_image=_UpperCAmelCase		)

							lowercase__					      =		pipe(
							    **_UpperCAmelCase				, return_dict=_UpperCAmelCase				, )[0]

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

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

							lowercase__					      =		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   lowerCamelCase__					(self   :						str		)						->       List[Any]:




							"""simple docstring"""


							lowercase__					      =		"""cpu"""

							lowercase__					      =		self.get_dummy_components()

							lowercase__					      =		self.pipeline_class(**_UpperCAmelCase		)
							lowercase__					      =		pipe.to(_UpperCAmelCase		)

							pipe.set_progress_bar_config(disable=_UpperCAmelCase		)

							lowercase__					      =		self.get_dummy_inputs(_UpperCAmelCase				, pil_image=_UpperCAmelCase		)

							lowercase__					      =		pipe(**_UpperCAmelCase		)
							lowercase__					      =		output.images

							lowercase__					      =		self.get_dummy_inputs(_UpperCAmelCase				, pil_image=_UpperCAmelCase		)

							lowercase__					      =		pipe(
							    **_UpperCAmelCase				, return_dict=_UpperCAmelCase				, )[0]

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

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

							lowercase__					      =		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   lowerCamelCase__					(self   :						List[Any]		)						->       List[Any]:




							"""simple docstring"""


							lowercase__					      =		"""cpu"""

							lowercase__					      =		self.get_dummy_components()

							lowercase__					      =		self.pipeline_class(**_UpperCAmelCase		)
							lowercase__					      =		pipe.to(_UpperCAmelCase		)

							pipe.set_progress_bar_config(disable=_UpperCAmelCase		)

							lowercase__					      =		self.get_dummy_inputs(_UpperCAmelCase				, pil_image=_UpperCAmelCase		)
							lowercase__					      =		[
							    pipeline_inputs["""image"""],
							    pipeline_inputs["""image"""],
							]

							lowercase__					      =		pipe(**_UpperCAmelCase		)
							lowercase__					      =		output.images

							lowercase__					      =		self.get_dummy_inputs(_UpperCAmelCase				, pil_image=_UpperCAmelCase		)
							lowercase__					      =		[
							    tuple_pipeline_inputs["""image"""],
							    tuple_pipeline_inputs["""image"""],
							]

							lowercase__					      =		pipe(
							    **_UpperCAmelCase				, return_dict=_UpperCAmelCase				, )[0]

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

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

							lowercase__					      =		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   lowerCamelCase__					(self   :						Tuple		)						->       Union[str, Any]:




							"""simple docstring"""


							lowercase__					      =		torch.device("""cpu"""		)



							class 							A   :
												'''simple docstring'''






												A__      =			1






							lowercase__					      =		self.get_dummy_components()

							lowercase__					      =		self.pipeline_class(**_UpperCAmelCase		)
							lowercase__					      =		pipe.to(_UpperCAmelCase		)

							pipe.set_progress_bar_config(disable=_UpperCAmelCase		)

							lowercase__					      =		torch.Generator(device=_UpperCAmelCase		).manual_seed(0		)
							lowercase__					      =		pipe.decoder.dtype
							lowercase__					      =		1

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

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

							lowercase__					      =		self.get_dummy_inputs(_UpperCAmelCase				, pil_image=_UpperCAmelCase		)

							lowercase__					      =		pipe(
							    **_UpperCAmelCase				, decoder_latents=_UpperCAmelCase				, super_res_latents=_UpperCAmelCase		).images

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

							lowercase__					      =		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   lowerCamelCase__					(self   :						str		)						->       Dict:




							"""simple docstring"""


							lowercase__					      =		torch_device == """cpu"""

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

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

					@skip_mps
					def   lowerCamelCase__					(self   :						Tuple		)						->       List[Any]:




							"""simple docstring"""


							lowercase__					      =		torch_device == """cpu"""
							lowercase__					      =		True
							lowercase__					      =		[
							    """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   lowerCamelCase__					(self   :						List[str]		)						->       Tuple:




							"""simple docstring"""


							lowercase__					      =		[
							    """decoder_num_inference_steps""",
							    """super_res_num_inference_steps""",
							]

							if torch_device == "mps":
									# TODO: MPS errors with larger batch sizes
									lowercase__					      =		[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   lowerCamelCase__					(self   :						Tuple		)						->       str:




							"""simple docstring"""


							return super().test_dict_tuple_outputs_equivalent()

					@skip_mps
					def   lowerCamelCase__					(self   :						int		)						->       Optional[Any]:




							"""simple docstring"""


							return super().test_save_load_local()


					@skip_mps
					def   lowerCamelCase__					(self   :						Optional[int]		)						->       List[Any]:




							"""simple docstring"""


							return super().test_save_load_optional_components()






@slow
@require_torch_gpu
class 							A   (	unittest.TestCase							):
					'''simple docstring'''






					def   lowerCamelCase__					(self   :						Optional[Any]		)						->       Dict:




							"""simple docstring"""


							super().tearDown()
							gc.collect()
							torch.cuda.empty_cache()


					def   lowerCamelCase__					(self   :						List[str]		)						->       Union[str, Any]:




							"""simple docstring"""


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

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

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

							lowercase__					      =		output.images[0]

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

							assert_mean_pixel_difference(_UpperCAmelCase				, _UpperCAmelCase				, 15		)





from __future__ import annotations

from collections import deque



class 							A   :
					'''simple docstring'''






					def __init__(self   :						Any				, _UpperCAmelCase   :						list[str]		)						->       Optional[int]:




							"""simple docstring"""


							lowercase__					      =		[]
							self.adlist.append(
							    {"""value""": """""", """next_states""": [], """fail_state""": 0, """output""": []}		)

							for keyword in keywords:
									self.add_keyword(_UpperCAmelCase		)
							self.set_fail_transitions()

					def   lowerCamelCase__					(self   :						List[str]				, _UpperCAmelCase   :						int				, _UpperCAmelCase   :						str		)						->       int | None:




							"""simple docstring"""


							for state in self.adlist[current_state]["next_states"]:
									if char == self.adlist[state]["value"]:
											return state
							return None

					def   lowerCamelCase__					(self   :						List[Any]				, _UpperCAmelCase   :						str		)						->       None:




							"""simple docstring"""


							lowercase__					      =		0
							for character in keyword:
									lowercase__					      =		self.find_next_state(_UpperCAmelCase				, _UpperCAmelCase		)
									if next_state is None:
											self.adlist.append(
											    {
											        """value""": character,
											        """next_states""": [],
											        """fail_state""": 0,
											        """output""": [],
											    }		)
											self.adlist[current_state]["next_states"].append(len(self.adlist		) - 1		)
											lowercase__					      =		len(self.adlist		) - 1
									else:
											lowercase__					      =		next_state
							self.adlist[current_state]["output"].append(_UpperCAmelCase		)

					def   lowerCamelCase__					(self   :						Any		)						->       None:




							"""simple docstring"""


							lowercase__					      =		deque()
							for node in self.adlist[0]["next_states"]:
									q.append(_UpperCAmelCase		)
									lowercase__					      =		0
							while q:
									lowercase__					      =		q.popleft()
									for child in self.adlist[r]["next_states"]:
											q.append(_UpperCAmelCase		)
											lowercase__					      =		self.adlist[r]["""fail_state"""]
											while (
											    self.find_next_state(_UpperCAmelCase				, self.adlist[child]["""value"""]		) is None
											    and state != 0
											):
													lowercase__					      =		self.adlist[state]["""fail_state"""]
											lowercase__					      =		self.find_next_state(
											    _UpperCAmelCase				, self.adlist[child]["""value"""]		)
											if self.adlist[child]["fail_state"] is None:
													lowercase__					      =		0
											lowercase__					      =		(
											    self.adlist[child]["""output"""]
											    + self.adlist[self.adlist[child]["""fail_state"""]]["""output"""]
											)


					def   lowerCamelCase__					(self   :						List[str]				, _UpperCAmelCase   :						str		)						->       dict[str, list[int]]:




							"""simple docstring"""


							lowercase__					      =		{}  # returns a dict with keywords and list of its occurrences
							lowercase__					      =		0
							for i in range(len(_UpperCAmelCase		)		):
									while (
									    self.find_next_state(_UpperCAmelCase				, string[i]		) is None
									    and current_state != 0
									):
											lowercase__					      =		self.adlist[current_state]["""fail_state"""]
									lowercase__					      =		self.find_next_state(_UpperCAmelCase				, string[i]		)
									if next_state is None:
											lowercase__					      =		0
									else:
											lowercase__					      =		next_state
											for key in self.adlist[current_state]["output"]:
													if key not in result:
															lowercase__					      =		[]
													result[key].append(i - len(_UpperCAmelCase		) + 1		)
							return result


if __name__ == "__main__":
	import doctest

	doctest.testmod()

