





import gc
import unittest

import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    DDPMScheduler,
    PriorTransformer,
    StableUnCLIPPipeline,
    UNetaDConditionModel,
)
from diffusers.pipelines.stable_diffusion.stable_unclip_image_normalizer import StableUnCLIPImageNormalizer
from diffusers.utils.testing_utils import enable_full_determinism, load_numpy, require_torch_gpu, slow, torch_device

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import (
    PipelineKarrasSchedulerTesterMixin,
    PipelineLatentTesterMixin,
    PipelineTesterMixin,
    assert_mean_pixel_difference,
)


enable_full_determinism()




class   UpperCamelCase_       (     _A				,  _A				,  _A				,  unittest.TestCase    ):
							lowerCamelCase_           =   StableUnCLIPPipeline
							lowerCamelCase_           =   TEXT_TO_IMAGE_PARAMS
							lowerCamelCase_           =   TEXT_TO_IMAGE_BATCH_PARAMS
							lowerCamelCase_           =   TEXT_TO_IMAGE_IMAGE_PARAMS
							lowerCamelCase_           =   TEXT_TO_IMAGE_IMAGE_PARAMS

							# TODO(will) Expected attn_bias.stride(1) == 0 to be true, but got false
							lowerCamelCase_           =   False


							def      _snake_case  (							self    :Optional[int]		)   ->   List[Any]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      32
									SCREAMING_SNAKE_CASE__			=      embedder_hidden_size

									# prior components

									torch.manual_seed(0		)
									SCREAMING_SNAKE_CASE__			=      CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""		)

									torch.manual_seed(0		)
									SCREAMING_SNAKE_CASE__			=      CLIPTextModelWithProjection(
									    CLIPTextConfig(
									        bos_token_id=0							,				eos_token_id=2							,				hidden_size=UpperCamelCase__							,				projection_dim=UpperCamelCase__							,				intermediate_size=37							,				layer_norm_eps=1E-05							,				num_attention_heads=4							,				num_hidden_layers=5							,				pad_token_id=1							,				vocab_size=1000							,				)		)

									torch.manual_seed(0		)
									SCREAMING_SNAKE_CASE__			=      PriorTransformer(
									    num_attention_heads=2							,				attention_head_dim=12							,				embedding_dim=UpperCamelCase__							,				num_layers=1							,				)

									torch.manual_seed(0		)
									SCREAMING_SNAKE_CASE__			=      DDPMScheduler(
									    variance_type="""fixed_small_log"""							,				prediction_type="""sample"""							,				num_train_timesteps=1000							,				clip_sample=UpperCamelCase__							,				clip_sample_range=5.0							,				beta_schedule="""squaredcos_cap_v2"""							,				)

									# regular denoising components

									torch.manual_seed(0		)
									SCREAMING_SNAKE_CASE__			=      StableUnCLIPImageNormalizer(embedding_dim=UpperCamelCase__		)
									SCREAMING_SNAKE_CASE__			=      DDPMScheduler(beta_schedule="""squaredcos_cap_v2"""		)

									torch.manual_seed(0		)
									SCREAMING_SNAKE_CASE__			=      CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""		)

									torch.manual_seed(0		)
									SCREAMING_SNAKE_CASE__			=      CLIPTextModel(
									    CLIPTextConfig(
									        bos_token_id=0							,				eos_token_id=2							,				hidden_size=UpperCamelCase__							,				projection_dim=32							,				intermediate_size=37							,				layer_norm_eps=1E-05							,				num_attention_heads=4							,				num_hidden_layers=5							,				pad_token_id=1							,				vocab_size=1000							,				)		)

									torch.manual_seed(0		)
									SCREAMING_SNAKE_CASE__			=      UNetaDConditionModel(
									    sample_size=32							,				in_channels=4							,				out_channels=4							,				down_block_types=("""CrossAttnDownBlock2D""", """DownBlock2D""")							,				up_block_types=("""UpBlock2D""", """CrossAttnUpBlock2D""")							,				block_out_channels=(32, 64)							,				attention_head_dim=(2, 4)							,				class_embed_type="""projection"""							,				projection_class_embeddings_input_dim=embedder_projection_dim * 2							,				cross_attention_dim=UpperCamelCase__							,				layers_per_block=1							,				upcast_attention=UpperCamelCase__							,				use_linear_projection=UpperCamelCase__							,				)

									torch.manual_seed(0		)
									SCREAMING_SNAKE_CASE__			=      DDIMScheduler(
									    beta_schedule="""scaled_linear"""							,				beta_start=0.0_0_0_8_5							,				beta_end=0.0_1_2							,				prediction_type="""v_prediction"""							,				set_alpha_to_one=UpperCamelCase__							,				steps_offset=1							,				)

									torch.manual_seed(0		)
									SCREAMING_SNAKE_CASE__			=      AutoencoderKL()

									SCREAMING_SNAKE_CASE__			=      {
									    # prior components
									    """prior_tokenizer""": prior_tokenizer,
									    """prior_text_encoder""": prior_text_encoder,
									    """prior""": prior,
									    """prior_scheduler""": prior_scheduler,
									    # image noising components
									    """image_normalizer""": image_normalizer,
									    """image_noising_scheduler""": image_noising_scheduler,
									    # regular denoising components
									    """tokenizer""": tokenizer,
									    """text_encoder""": text_encoder,
									    """unet""": unet,
									    """scheduler""": scheduler,
									    """vae""": vae,
									}

									return components


							def      _snake_case  (							self    :List[Any]							,				__A    :int							,				__A    :Any=0		)   ->   Tuple:



									"""simple docstring"""



									if str(UpperCamelCase__		).startswith("""mps"""		):
											SCREAMING_SNAKE_CASE__			=      torch.manual_seed(UpperCamelCase__		)
									else:
											SCREAMING_SNAKE_CASE__			=      torch.Generator(device=UpperCamelCase__		).manual_seed(UpperCamelCase__		)
									SCREAMING_SNAKE_CASE__			=      {
									    """prompt""": """A painting of a squirrel eating a burger""",
									    """generator""": generator,
									    """num_inference_steps""": 2,
									    """prior_num_inference_steps""": 2,
									    """output_type""": """numpy""",
									}
									return inputs


							def      _snake_case  (							self    :Optional[Any]		)   ->   Any:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      torch_device == """cpu"""

									self._test_attention_slicing_forward_pass(test_max_difference=UpperCamelCase__		)


							def      _snake_case  (							self    :int		)   ->   int:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      torch_device in ["""cpu""", """mps"""]

									self._test_inference_batch_single_identical(test_max_difference=UpperCamelCase__		)




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


							def      _snake_case  (							self    :int		)   ->   int:



									"""simple docstring"""



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


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



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      load_numpy(
									    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_l_anime_turtle_fp16.npy"""		)

									SCREAMING_SNAKE_CASE__			=      StableUnCLIPPipeline.from_pretrained("""fusing/stable-unclip-2-1-l"""							,				torch_dtype=torch.floataa		)
									pipe.to(UpperCamelCase__		)
									pipe.set_progress_bar_config(disable=UpperCamelCase__		)
									# stable unclip will oom when integration tests are run on a V100,
									# so turn on memory savings
									pipe.enable_attention_slicing()
									pipe.enable_sequential_cpu_offload()

									SCREAMING_SNAKE_CASE__			=      torch.Generator(device="""cpu"""		).manual_seed(0		)
									SCREAMING_SNAKE_CASE__			=      pipe("""anime turle"""							,				generator=UpperCamelCase__							,				output_type="""np"""		)

									SCREAMING_SNAKE_CASE__			=      output.images[0]

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

									assert_mean_pixel_difference(UpperCamelCase__							,				UpperCamelCase__		)


							def      _snake_case  (							self    :List[Any]		)   ->   Optional[Any]:



									"""simple docstring"""



									torch.cuda.empty_cache()
									torch.cuda.reset_max_memory_allocated()
									torch.cuda.reset_peak_memory_stats()

									SCREAMING_SNAKE_CASE__			=      StableUnCLIPPipeline.from_pretrained("""fusing/stable-unclip-2-1-l"""							,				torch_dtype=torch.floataa		)
									SCREAMING_SNAKE_CASE__			=      pipe.to(UpperCamelCase__		)
									pipe.set_progress_bar_config(disable=UpperCamelCase__		)
									pipe.enable_attention_slicing()
									pipe.enable_sequential_cpu_offload()

									SCREAMING_SNAKE_CASE__			=      pipe(
									    """anime turtle"""							,				prior_num_inference_steps=2							,				num_inference_steps=2							,				output_type="""np"""							,				)

									SCREAMING_SNAKE_CASE__			=      torch.cuda.max_memory_allocated()
									# make sure that less than 7 GB is allocated
									assert mem_bytes < 7 * 10**9






def 				SCREAMING_SNAKE_CASE__							(       UpperCamelCase__:							list[list[float]]    ):
		SCREAMING_SNAKE_CASE__			=      []
		for data in source_data:
				for i, el in enumerate(UpperCamelCase__    ):
						if len(UpperCamelCase__    ) < i + 1:
								data_lists.append([]    )
						data_lists[i].append(float(UpperCamelCase__    )    )
		return data_lists






def 				SCREAMING_SNAKE_CASE__							(       UpperCamelCase__:							list[list[float]]					,   UpperCamelCase__:							list[int]    ):
		SCREAMING_SNAKE_CASE__			=      []
		for dlist, weight in zip(UpperCamelCase__					,   UpperCamelCase__    ):
				SCREAMING_SNAKE_CASE__			=      min(UpperCamelCase__    )
				SCREAMING_SNAKE_CASE__			=      max(UpperCamelCase__    )

				SCREAMING_SNAKE_CASE__			=      []
				# for weight 0 score is 1 - actual score
				if weight == 0:
						for item in dlist:
								try:
										score.append(1 - ((item - mind) / (maxd - mind))    )
								except ZeroDivisionError:
										score.append(1    )

				elif weight == 1:
						for item in dlist:
								try:
										score.append((item - mind) / (maxd - mind)    )
								except ZeroDivisionError:
										score.append(0    )

        # weight not 0 or 1
				else:
						SCREAMING_SNAKE_CASE__			=      f'''Invalid weight of {weight:f} provided'''
						raise ValueError(UpperCamelCase__    )

				score_lists.append(UpperCamelCase__    )

		return score_lists






def 				SCREAMING_SNAKE_CASE__							(       UpperCamelCase__:							list[list[float]]    ):
		SCREAMING_SNAKE_CASE__			=      [0 for i in range(len(score_lists[0]    )    )]

		for slist in score_lists:
				for j, ele in enumerate(UpperCamelCase__    ):
						SCREAMING_SNAKE_CASE__			=      final_scores[j] + ele

		return final_scores






def 				SCREAMING_SNAKE_CASE__							(       UpperCamelCase__:							list[list[float]]					,   UpperCamelCase__:							list[int]    ):
		SCREAMING_SNAKE_CASE__			=      get_data(UpperCamelCase__    )
		SCREAMING_SNAKE_CASE__			=      calculate_each_score(UpperCamelCase__					,   UpperCamelCase__    )
		SCREAMING_SNAKE_CASE__			=      generate_final_scores(UpperCamelCase__    )

		# append scores to source data
		for i, ele in enumerate(UpperCamelCase__    ):
				source_data[i].append(UpperCamelCase__    )

		return source_data