


"""simple docstring"""






import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image

from diffusers import (
    DDIMScheduler,
    KandinskyVaaImgaImgPipeline,
    KandinskyVaaPriorPipeline,
    UNetaDConditionModel,
    VQModel,
)
from diffusers.utils import floats_tensor, load_image, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()


class 			lowerCamelCase_		(			lowercase       ,  unittest.TestCase   ):





				"""simple docstring"""







				_lowerCAmelCase      :      Optional[int]      							=					KandinskyVaaImgaImgPipeline
				_lowerCAmelCase      :      Dict      							=					["image_embeds", "negative_image_embeds", "image"]
				_lowerCAmelCase      :      Union[str, Any]      							=					[
				    "image_embeds",
				    "negative_image_embeds",
				    "image",
				]
				_lowerCAmelCase      :      str      							=					[
				    "generator",
				    "height",
				    "width",
				    "strength",
				    "guidance_scale",
				    "num_inference_steps",
				    "return_dict",
				    "guidance_scale",
				    "num_images_per_prompt",
				    "output_type",
				    "return_dict",
				]
				_lowerCAmelCase      :      Union[str, Any]      							=					False



				@property
				def   lowerCAmelCase__		(						self     ):
								return 32



				@property
				def   lowerCAmelCase__		(						self     ):
								return 32



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



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



				@property
				def   lowerCAmelCase__		(						self     ):
								return 100



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

								SCREAMING_SNAKE_CASE__     							=				{
								    "in_channels": 4,
								    # Out channels is double in channels because predicts mean and variance
								    "out_channels": 8,
								    "addition_embed_type": "image",
								    "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,
								    "encoder_hid_dim": self.text_embedder_hidden_size,
								    "encoder_hid_dim_type": "image_proj",
								    "cross_attention_dim": self.cross_attention_dim,
								    "attention_head_dim": 4,
								    "resnet_time_scale_shift": "scale_shift",
								    "class_embed_type": None,
								}

								SCREAMING_SNAKE_CASE__     							=				UNetaDConditionModel(**UpperCamelCase__     )
								return model



				@property
				def   lowerCAmelCase__		(						self     ):
								return {
								    "block_out_channels": [32, 64],
								    "down_block_types": ["DownEncoderBlock2D", "AttnDownEncoderBlock2D"],
								    "in_channels": 3,
								    "latent_channels": 4,
								    "layers_per_block": 1,
								    "norm_num_groups": 8,
								    "norm_type": "spatial",
								    "num_vq_embeddings": 12,
								    "out_channels": 3,
								    "up_block_types": [
								        "AttnUpDecoderBlock2D",
								        "UpDecoderBlock2D",
								    ],
								    "vq_embed_dim": 4,
								}



				@property
				def   lowerCAmelCase__		(						self     ):
								torch.manual_seed(0     )
								SCREAMING_SNAKE_CASE__     							=				VQModel(**self.dummy_movq_kwargs     )
								return model



				def   lowerCAmelCase__		(						self     ):
								SCREAMING_SNAKE_CASE__     							=				self.dummy_unet
								SCREAMING_SNAKE_CASE__     							=				self.dummy_movq

								SCREAMING_SNAKE_CASE__     							=				{
								    "num_train_timesteps": 1000,
								    "beta_schedule": "linear",
								    "beta_start": 0.00_085,
								    "beta_end": 0.012,
								    "clip_sample": False,
								    "set_alpha_to_one": False,
								    "steps_offset": 0,
								    "prediction_type": "epsilon",
								    "thresholding": False,
								}

								SCREAMING_SNAKE_CASE__     							=				DDIMScheduler(**UpperCamelCase__     )

								SCREAMING_SNAKE_CASE__     							=				{
								    "unet": unet,
								    "scheduler": scheduler,
								    "movq": movq,
								}

								return components



				def   lowerCAmelCase__		(						self       ,		UpperCAmelCase__       ,		UpperCAmelCase__=0     ):
								SCREAMING_SNAKE_CASE__     							=				floats_tensor((1, self.text_embedder_hidden_size)       ,		rng=random.Random(UpperCamelCase__     )     ).to(UpperCamelCase__     )
								SCREAMING_SNAKE_CASE__     							=				floats_tensor((1, self.text_embedder_hidden_size)       ,		rng=random.Random(seed + 1     )     ).to(
								    UpperCamelCase__     )
								# create init_image
								SCREAMING_SNAKE_CASE__     							=				floats_tensor((1, 3, 64, 64)       ,		rng=random.Random(UpperCamelCase__     )     ).to(UpperCamelCase__     )
								SCREAMING_SNAKE_CASE__     							=				image.cpu().permute(0       ,		2       ,		3       ,		1     )[0]
								SCREAMING_SNAKE_CASE__     							=				Image.fromarray(np.uinta(UpperCamelCase__     )     ).convert("RGB"     ).resize((256, 256)     )

								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__     							=				{
								    "image": init_image,
								    "image_embeds": image_embeds,
								    "negative_image_embeds": negative_image_embeds,
								    "generator": generator,
								    "height": 64,
								    "width": 64,
								    "num_inference_steps": 10,
								    "guidance_scale": 7.0,
								    "strength": 0.2,
								    "output_type": "np",
								}
								return inputs






				def   lowerCAmelCase__		(						self     ):
								SCREAMING_SNAKE_CASE__     							=				"cpu"

								SCREAMING_SNAKE_CASE__     							=				self.get_dummy_components()

								SCREAMING_SNAKE_CASE__     							=				self.pipeline_class(**UpperCamelCase__     )
								SCREAMING_SNAKE_CASE__     							=				pipe.to(UpperCamelCase__     )

								pipe.set_progress_bar_config(disable=UpperCamelCase__     )

								SCREAMING_SNAKE_CASE__     							=				pipe(**self.get_dummy_inputs(UpperCamelCase__     )     )
								SCREAMING_SNAKE_CASE__     							=				output.images

								SCREAMING_SNAKE_CASE__     							=				pipe(
								    **self.get_dummy_inputs(UpperCamelCase__     )       ,		return_dict=UpperCamelCase__       ,		)[0]

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

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

								SCREAMING_SNAKE_CASE__     							=				np.array(
								    [0.6_199_778, 0.63_984_406, 0.46_145_785, 0.62_944_984, 0.5_622_215, 0.47_306_132, 0.47_441_456, 0.4_607_606, 0.48_719_263]     )
								assert (
								    np.abs(image_slice.flatten() - expected_slice     ).max() < 1e-2
								), f''' expected_slice {expected_slice}, but got {image_slice.flatten()}'''
								assert (
								    np.abs(image_from_tuple_slice.flatten() - expected_slice     ).max() < 1e-2
								), f''' expected_slice {expected_slice}, but got {image_from_tuple_slice.flatten()}'''


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





				"""simple docstring"""







				def   lowerCAmelCase__		(						self     ):
								super().tearDown()
								gc.collect()
								torch.cuda.empty_cache()






				def   lowerCAmelCase__		(						self     ):
								SCREAMING_SNAKE_CASE__     							=				load_numpy(
								    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
								    "/kandinskyv22/kandinskyv22_img2img_frog.npy"     )

								SCREAMING_SNAKE_CASE__     							=				load_image(
								    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main" "/kandinsky/cat.png"     )
								SCREAMING_SNAKE_CASE__     							=				"A red cartoon frog, 4k"

								SCREAMING_SNAKE_CASE__     							=				KandinskyVaaPriorPipeline.from_pretrained(
								    "kandinsky-community/kandinsky-2-2-prior"       ,		torch_dtype=torch.floataa     )
								pipe_prior.to(UpperCamelCase__     )

								SCREAMING_SNAKE_CASE__     							=				KandinskyVaaImgaImgPipeline.from_pretrained(
								    "kandinsky-community/kandinsky-2-2-decoder"       ,		torch_dtype=torch.floataa     )
								SCREAMING_SNAKE_CASE__     							=				pipeline.to(UpperCamelCase__     )

								pipeline.set_progress_bar_config(disable=UpperCamelCase__     )

								SCREAMING_SNAKE_CASE__     							=				torch.Generator(device="cpu"     ).manual_seed(0     )
								SCREAMING_SNAKE_CASE__							,      SCREAMING_SNAKE_CASE__     							=				pipe_prior(
								    UpperCamelCase__       ,		generator=UpperCamelCase__       ,		num_inference_steps=5       ,		negative_prompt=""       ,		).to_tuple()

								SCREAMING_SNAKE_CASE__     							=				pipeline(
								    image=UpperCamelCase__       ,		image_embeds=UpperCamelCase__       ,		negative_image_embeds=UpperCamelCase__       ,		generator=UpperCamelCase__       ,		num_inference_steps=100       ,		height=768       ,		width=768       ,		strength=0.2       ,		output_type="np"       ,		)

								SCREAMING_SNAKE_CASE__     							=				output.images[0]

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

								assert_mean_pixel_difference(UpperCamelCase__       ,		UpperCamelCase__     )




"""simple docstring"""






import os

from bleurt import score  # From: git+https://github.com/google-research/bleurt.git

import datasets


_lowerCamelCase     =  datasets.logging.get_logger(__name__)


_lowerCamelCase     =  '\\n@inproceedings{bleurt,\n  title={BLEURT: Learning Robust Metrics for Text Generation},\n  author={Thibault Sellam and Dipanjan Das and Ankur P. Parikh},\n  booktitle={ACL},\n  year={2020},\n  url={https://arxiv.org/abs/2004.04696}\n}\n'

_lowerCamelCase     =  '\\nBLEURT a learnt evaluation metric for Natural Language Generation. It is built using multiple phases of transfer learning starting from a pretrained BERT model (Devlin et al. 2018)\nand then employing another pre-training phrase using synthetic data. Finally it is trained on WMT human annotations. You may run BLEURT out-of-the-box or fine-tune\nit for your specific application (the latter is expected to perform better).\n\nSee the project\'s README at https://github.com/google-research/bleurt#readme for more information.\n'

_lowerCamelCase     =  '\nBLEURT score.\n\nArgs:\n    `predictions` (list of str): prediction/candidate sentences\n    `references` (list of str): reference sentences\n    `checkpoint` BLEURT checkpoint. Will default to BLEURT-tiny if None.\n\nReturns:\n    \'scores\': List of scores.\nExamples:\n\n    >>> predictions = ["hello there", "general kenobi"]\n    >>> references = ["hello there", "general kenobi"]\n    >>> bleurt = datasets.load_metric("bleurt")\n    >>> results = bleurt.compute(predictions=predictions, references=references)\n    >>> print([round(v, 2) for v in results["scores"]])\n    [1.03, 1.04]\n'

_lowerCamelCase     =  {
    'bleurt-tiny-128': 'https://storage.googleapis.com/bleurt-oss/bleurt-tiny-128.zip',
    'bleurt-tiny-512': 'https://storage.googleapis.com/bleurt-oss/bleurt-tiny-512.zip',
    'bleurt-base-128': 'https://storage.googleapis.com/bleurt-oss/bleurt-base-128.zip',
    'bleurt-base-512': 'https://storage.googleapis.com/bleurt-oss/bleurt-base-512.zip',
    'bleurt-large-128': 'https://storage.googleapis.com/bleurt-oss/bleurt-large-128.zip',
    'bleurt-large-512': 'https://storage.googleapis.com/bleurt-oss/bleurt-large-512.zip',
    'BLEURT-20-D3': 'https://storage.googleapis.com/bleurt-oss-21/BLEURT-20-D3.zip',
    'BLEURT-20-D6': 'https://storage.googleapis.com/bleurt-oss-21/BLEURT-20-D6.zip',
    'BLEURT-20-D12': 'https://storage.googleapis.com/bleurt-oss-21/BLEURT-20-D12.zip',
    'BLEURT-20': 'https://storage.googleapis.com/bleurt-oss-21/BLEURT-20.zip',
}


@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION       ,  _KWARGS_DESCRIPTION   )
class 			lowerCamelCase_		(			datasets.Metric   ):





				"""simple docstring"""







				def   lowerCAmelCase__		(						self     ):
								return datasets.MetricInfo(
								    description=_DESCRIPTION       ,		citation=_CITATION       ,		homepage="https://github.com/google-research/bleurt"       ,		inputs_description=_KWARGS_DESCRIPTION       ,		features=datasets.Features(
								        {
								            "predictions": datasets.Value("string"       ,		id="sequence"     ),
								            "references": datasets.Value("string"       ,		id="sequence"     ),
								        }     )       ,		codebase_urls=["https://github.com/google-research/bleurt"]       ,		reference_urls=["https://github.com/google-research/bleurt", "https://arxiv.org/abs/2004.04696"]       ,		)



				def   lowerCAmelCase__		(						self       ,		UpperCAmelCase__     ):
								# check that config name specifies a valid BLEURT model
								if self.config_name == "default":
												logger.warning(
												    "Using default BLEURT-Base checkpoint for sequence maximum length 128. "
												    "You can use a bigger model for better results with e.g.: datasets.load_metric('bleurt', 'bleurt-large-512')."     )
												SCREAMING_SNAKE_CASE__     							=				"bleurt-base-128"

								if self.config_name.lower() in CHECKPOINT_URLS:
												SCREAMING_SNAKE_CASE__     							=				self.config_name.lower()

								elif self.config_name.upper() in CHECKPOINT_URLS:
												SCREAMING_SNAKE_CASE__     							=				self.config_name.upper()

								else:
												raise KeyError(
												    f'''{self.config_name} model not found. You should supply the name of a model checkpoint for bleurt in {CHECKPOINT_URLS.keys()}'''     )

								# download the model checkpoint specified by self.config_name and set up the scorer
								SCREAMING_SNAKE_CASE__     							=				dl_manager.download_and_extract(CHECKPOINT_URLS[checkpoint_name]     )
								SCREAMING_SNAKE_CASE__     							=				score.BleurtScorer(os.path.join(UpperCAmelCase__       ,		UpperCAmelCase__     )     )






				def   lowerCAmelCase__		(						self       ,		UpperCAmelCase__       ,		UpperCAmelCase__     ):
								SCREAMING_SNAKE_CASE__     							=				self.scorer.score(references=UpperCAmelCase__       ,		candidates=UpperCAmelCase__     )
								return {"scores": scores}
