



import gc
import random
import unittest

import numpy as np
import torch
from transformers import XLMRobertaTokenizerFast

from diffusers import DDIMScheduler, KandinskyPipeline, KandinskyPriorPipeline, UNetaDConditionModel, VQModel
from diffusers.pipelines.kandinsky.text_encoder import MCLIPConfig, MultilingualCLIP
from diffusers.utils import floats_tensor, 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__  (					_A       ,     unittest.TestCase):



 """simple docstring"""

 a__				:   List[str]   =     KandinskyPipeline
 a__				:   Tuple   =     [
     "prompt",
     "image_embeds",
     "negative_image_embeds",
 ]
 a__				:   Dict   =     ["prompt", "negative_prompt", "image_embeds", "negative_image_embeds"]
 a__				:   List[Any]   =     [
     "generator",
     "height",
     "width",
     "latents",
     "guidance_scale",
     "negative_prompt",
     "num_inference_steps",
     "return_dict",
     "guidance_scale",
     "num_images_per_prompt",
     "output_type",
     "return_dict",
 ]
 a__				:   List[str]   =     False

 @property
 def 						snake_case_			(						self				:				int			)      ->  int:
    return 32

 @property
 def 						snake_case_			(						self				:				int			)      ->  List[str]:
    return 32

 @property
 def 						snake_case_			(						self				:				Tuple			)      ->  Any:
    return self.time_input_dim

 @property
 def 						snake_case_			(						self				:				Any			)      ->  Optional[Any]:
    return self.time_input_dim * 4

 @property
 def 						snake_case_			(						self				:				Tuple			)      ->  Tuple:
    return 1_00

 @property
 def 						snake_case_			(						self				:				str			)      ->  Optional[int]:
    _A						     =					XLMRobertaTokenizerFast.from_pretrained('''YiYiXu/tiny-random-mclip-base'''			)
    return tokenizer

 @property
 def 						snake_case_			(						self				:				Optional[int]			)      ->  Any:
    torch.manual_seed(0			)
    _A						     =					MCLIPConfig(
        numDims=self.cross_attention_dim ,	transformerDimensions=self.text_embedder_hidden_size ,	hidden_size=self.text_embedder_hidden_size ,	intermediate_size=37 ,	num_attention_heads=4 ,	num_hidden_layers=5 ,	vocab_size=10_05 ,	)

    _A						     =					MultilingualCLIP(__lowerCAmelCase			)
    _A						     =					text_encoder.eval()

    return text_encoder

 @property
 def 						snake_case_			(						self				:				Union[str, Any]			)      ->  Dict:
    torch.manual_seed(0			)

    _A						     =					{
        '''in_channels''': 4,
        # Out channels is double in channels because predicts mean and variance
        '''out_channels''': 8,
        '''addition_embed_type''': '''text_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''': '''text_image_proj''',
        '''cross_attention_dim''': self.cross_attention_dim,
        '''attention_head_dim''': 4,
        '''resnet_time_scale_shift''': '''scale_shift''',
        '''class_embed_type''': None,
    }

    _A						     =					UNetaDConditionModel(**__lowerCAmelCase			)
    return model

 @property
 def 						snake_case_			(						self				:				Dict			)      ->  str:
    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 						snake_case_			(						self				:				Optional[Any]			)      ->  Any:
    torch.manual_seed(0			)
    _A						     =					VQModel(**self.dummy_movq_kwargs			)
    return model

 def 						snake_case_			(						self				:				Union[str, Any]			)      ->  Any:
    _A						     =					self.dummy_text_encoder
    _A						     =					self.dummy_tokenizer
    _A						     =					self.dummy_unet
    _A						     =					self.dummy_movq

    _A						     =					DDIMScheduler(
        num_train_timesteps=10_00 ,	beta_schedule='''linear''' ,	beta_start=0.0_0085 ,	beta_end=0.012 ,	clip_sample=__lowerCAmelCase ,	set_alpha_to_one=__lowerCAmelCase ,	steps_offset=1 ,	prediction_type='''epsilon''' ,	thresholding=__lowerCAmelCase ,	)

    _A						     =					{
        '''text_encoder''': text_encoder,
        '''tokenizer''': tokenizer,
        '''unet''': unet,
        '''scheduler''': scheduler,
        '''movq''': movq,
    }
    return components

 def 						snake_case_			(						self				:				List[Any] ,	__lowerCAmelCase				:				List[str] ,	__lowerCAmelCase				:				str=0			)      ->  int:
    _A						     =					floats_tensor((1, self.cross_attention_dim) ,	rng=random.Random(__lowerCAmelCase			)			).to(__lowerCAmelCase			)
    _A						     =					floats_tensor((1, self.cross_attention_dim) ,	rng=random.Random(seed + 1			)			).to(__lowerCAmelCase			)
    if str(__lowerCAmelCase			).startswith('''mps'''			):
       _A						     =					torch.manual_seed(__lowerCAmelCase			)
    else:
       _A						     =					torch.Generator(device=__lowerCAmelCase			).manual_seed(__lowerCAmelCase			)
    _A						     =					{
        '''prompt''': '''horse''',
        '''image_embeds''': image_embeds,
        '''negative_image_embeds''': negative_image_embeds,
        '''generator''': generator,
        '''height''': 64,
        '''width''': 64,
        '''guidance_scale''': 4.0,
        '''num_inference_steps''': 2,
        '''output_type''': '''np''',
    }
    return inputs

 def 						snake_case_			(						self				:				Dict			)      ->  Optional[int]:
    _A						     =					'''cpu'''

    _A						     =					self.get_dummy_components()

    _A						     =					self.pipeline_class(**__lowerCAmelCase			)
    _A						     =					pipe.to(__lowerCAmelCase			)

    pipe.set_progress_bar_config(disable=__lowerCAmelCase			)

    _A						     =					pipe(**self.get_dummy_inputs(__lowerCAmelCase			)			)
    _A						     =					output.images

    _A						     =					pipe(
        **self.get_dummy_inputs(__lowerCAmelCase			) ,	return_dict=__lowerCAmelCase ,	)[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.32_8663, 1.0, 0.2321_6873, 1.0, 0.9271_7564, 0.463_9046, 0.9689_4777, 0.3171_3378, 0.629_3953]			)

    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 						snake_case_			(						self				:				Union[str, Any]			)      ->  List[Any]:
    # clean up the VRAM after each test
    super().tearDown()
    gc.collect()
    torch.cuda.empty_cache()

 def 						snake_case_			(						self				:				Tuple			)      ->  str:
    _A						     =					load_numpy(
        '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
        '''/kandinsky/kandinsky_text2img_cat_fp16.npy'''			)

    _A						     =					KandinskyPriorPipeline.from_pretrained(
        '''kandinsky-community/kandinsky-2-1-prior''' ,	torch_dtype=torch.floataa			)
    pipe_prior.to(__lowerCAmelCase			)

    _A						     =					KandinskyPipeline.from_pretrained('''kandinsky-community/kandinsky-2-1''' ,	torch_dtype=torch.floataa			)
    _A						     =					pipeline.to(__lowerCAmelCase			)
    pipeline.set_progress_bar_config(disable=__lowerCAmelCase			)

    _A						     =					'''red cat, 4k photo'''

    _A						     =					torch.Generator(device='''cuda'''			).manual_seed(0			)
    _A     ,						_A						     =					pipe_prior(
        __lowerCAmelCase ,	generator=__lowerCAmelCase ,	num_inference_steps=5 ,	negative_prompt='''''' ,	).to_tuple()

    _A						     =					torch.Generator(device='''cuda'''			).manual_seed(0			)
    _A						     =					pipeline(
        __lowerCAmelCase ,	image_embeds=__lowerCAmelCase ,	negative_image_embeds=__lowerCAmelCase ,	generator=__lowerCAmelCase ,	num_inference_steps=1_00 ,	output_type='''np''' ,	)

    _A						     =					output.images[0]

    assert image.shape == (5_12, 5_12, 3)

    assert_mean_pixel_difference(__lowerCAmelCase ,	__lowerCAmelCase			)








from typing import List

from packaging import version
from sklearn.metrics import fa_score

import datasets
from datasets.config import PY_VERSION


if PY_VERSION < version.parse("""3.8"""):
   import importlib_metadata
else:
   import importlib.metadata as importlib_metadata


# TODO(Patrick/Anton)
UpperCAmelCase_		   =				"""\
"""

UpperCAmelCase_		   =				"""\
    XTREME-S is a benchmark to evaluate universal cross-lingual speech representations in many languages.
    XTREME-S covers four task families: speech recognition, classification, speech-to-text translation and retrieval.
"""

UpperCAmelCase_		   =				"""
Compute XTREME-S evaluation metric associated to each XTREME-S dataset.
Args:
    predictions: list of predictions to score.
        Each translation should be tokenized into a list of tokens.
    references: list of lists of references for each translation.
        Each reference should be tokenized into a list of tokens.
    bleu_kwargs: optional dict of keywords to be passed when computing 'bleu'.
        Keywords include Dict can be one of 'smooth_method', 'smooth_value', 'force', 'lowercase',
        'tokenize', 'use_effective_order'.
    wer_kwargs: optional dict of keywords to be passed when computing 'wer' and 'cer'.
        Keywords include 'concatenate_texts'.
Returns: depending on the XTREME-S task, one or several of:
    \"accuracy\": Accuracy - for 'fleurs-lang_id', 'minds14'
    \"f1\": F1 score - for 'minds14'
    \"wer\": Word error rate - for 'mls', 'fleurs-asr', 'voxpopuli', 'babel'
    \"cer\": Character error rate - for 'mls', 'fleurs-asr', 'voxpopuli', 'babel'
    \"bleu\": BLEU score according to the `sacrebleu` metric - for 'covost2'
Examples:

    >>> xtreme_s_metric = datasets.load_metric('xtreme_s', 'mls')  # 'mls', 'voxpopuli', 'fleurs-asr' or 'babel'
    >>> references = [\"it is sunny here\", \"paper and pen are essentials\"]
    >>> predictions = [\"it's sunny\", \"paper pen are essential\"]
    >>> results = xtreme_s_metric.compute(predictions=predictions, references=references)
    >>> print({k: round(v, 2) for k, v in results.items()})
    {'wer': 0.56, 'cer': 0.27}

    >>> xtreme_s_metric = datasets.load_metric('xtreme_s', 'covost2')
    >>> references = [\"bonjour paris\", \"il est necessaire de faire du sport de temps en temp\"]
    >>> predictions = [\"bonjour paris\", \"il est important de faire du sport souvent\"]
    >>> results = xtreme_s_metric.compute(predictions=predictions, references=references)
    >>> print({k: round(v, 2) for k, v in results.items()})
    {'bleu': 31.65}

    >>> xtreme_s_metric = datasets.load_metric('xtreme_s', 'fleurs-lang_id')
    >>> references = [0, 1, 0, 0, 1]
    >>> predictions = [0, 1, 1, 0, 0]
    >>> results = xtreme_s_metric.compute(predictions=predictions, references=references)
    >>> print({k: round(v, 2) for k, v in results.items()})
    {'accuracy': 0.6}

    >>> xtreme_s_metric = datasets.load_metric('xtreme_s', 'minds14')
    >>> references = [0, 1, 0, 0, 1]
    >>> predictions = [0, 1, 1, 0, 0]
    >>> results = xtreme_s_metric.compute(predictions=predictions, references=references)
    >>> print({k: round(v, 2) for k, v in results.items()})
    {'f1': 0.58, 'accuracy': 0.6}
"""

UpperCAmelCase_		   =				["""fleurs-asr""", """mls""", """voxpopuli""", """babel""", """covost2""", """fleurs-lang_id""", """minds14"""]
UpperCAmelCase_		   =				""""""

try:
   from jiwer import transforms as tr

   UpperCAmelCase_		   =				True
except ImportError:
   UpperCAmelCase_		   =				False

if _jiwer_available and version.parse(importlib_metadata.version("""jiwer""")) < version.parse("""2.3.0"""):


   class   lowerCamelCase__  (					tr.AbstractTransform):



    """simple docstring"""

    def __init__(						self				:				str ,	__lowerCAmelCase				:				str = " "			)      ->  Any:
       _A						     =					sentence_delimiter

    def 						snake_case_			(						self				:				List[str] ,	__lowerCAmelCase				:				str			)      ->  Optional[int]:
       return list(__lowerCAmelCase			)

    def 						snake_case_			(						self				:				Union[str, Any] ,	__lowerCAmelCase				:				List[str]			)      ->  Optional[Any]:
       _A						     =					[]
       for sent_idx, sentence in enumerate(__lowerCAmelCase			):
          chars.extend(self.process_string(__lowerCAmelCase			)			)
          if self.sentence_delimiter is not None and self.sentence_delimiter != "" and sent_idx < len(__lowerCAmelCase			) - 1:
             chars.append(self.sentence_delimiter			)
       return chars
   UpperCAmelCase_		   =				tr.Compose(
       [tr.RemoveMultipleSpaces(), tr.Strip(), SentencesToListOfCharacters(SENTENCE_DELIMITER)]
   )
elif _jiwer_available:
   UpperCAmelCase_		   =				tr.Compose(
       [
           tr.RemoveMultipleSpaces(),
           tr.Strip(),
           tr.ReduceToSingleSentence(SENTENCE_DELIMITER),
           tr.ReduceToListOfListOfChars(),
       ]
   )
else:
   UpperCAmelCase_		   =				None

def   SCREAMING_SNAKE_CASE_       (  _snake_case  :Dict  ,      _snake_case  :Union[str, Any]							)     ->    int:
   return float((preds == labels).mean()							)

def   SCREAMING_SNAKE_CASE_       (  _snake_case  :int  ,      _snake_case  :Dict							)     ->    Optional[Any]:
   return {
       "f1": float(fa_score(y_true=_snake_case  ,      y_pred=_snake_case  ,      average='''macro'''							)							),
       "accuracy": simple_accuracy(_snake_case  ,      _snake_case							),
   }

def   SCREAMING_SNAKE_CASE_       (  _snake_case  :Union[str, Any]  ,      _snake_case  :Union[str, Any]  ,      _snake_case  :Tuple="exp"  ,      _snake_case  :Dict=None  ,      _snake_case  :List[str]=False  ,      _snake_case  :Union[str, Any]=False  ,      _snake_case  :List[Any]=None  ,      _snake_case  :Tuple=False  ,      )     ->    Tuple:
   # xtreme-s can only have one label
   _A						     =					[[label] for label in labels]
   _A						     =					list(_snake_case							)
   try:
      import sacrebleu as scb
   except ImportError:
      raise ValueError(
          '''sacrebleu has to be installed in order to apply the bleu metric for covost2.'''
          '''You can install it via `pip install sacrebleu`.'''							)

   if version.parse(scb.__version__							) < version.parse('''1.4.12'''							):
      raise ImportWarning(
          '''To use `sacrebleu`, the module `sacrebleu>=1.4.12` is required, and the current version of `sacrebleu` doesn\'t match this condition.\n'''
          '''You can install it with `pip install "sacrebleu>=1.4.12"`.'''							)

   _A						     =					len(labels[0]							)
   if any(len(_snake_case							) != references_per_prediction for refs in labels							):
      raise ValueError('''Sacrebleu requires the same number of references for each prediction'''							)
   _A						     =					[[refs[i] for refs in labels] for i in range(_snake_case							)]
   _A						     =					scb.corpus_bleu(
       _snake_case  ,      _snake_case  ,      smooth_method=_snake_case  ,      smooth_value=_snake_case  ,      force=_snake_case  ,      lowercase=_snake_case  ,      use_effective_order=_snake_case  ,      **({'''tokenize''': tokenize} if tokenize else {})  ,      )
   return {"bleu": output.score}

def   SCREAMING_SNAKE_CASE_       (  _snake_case  :Optional[Any]  ,      _snake_case  :Optional[int]  ,      _snake_case  :List[Any]  ,      _snake_case  :str							)     ->    Tuple:
   try:
      from jiwer import compute_measures
   except ImportError:
      raise ValueError(
          F'''jiwer has to be installed in order to apply the wer metric for {config_name}.'''
          '''You can install it via `pip install jiwer`.'''							)

   if concatenate_texts:
      _A						     =					compute_measures(_snake_case  ,      _snake_case							)['''wer''']

      _A						     =					compute_measures(_snake_case  ,      _snake_case  ,      truth_transform=_snake_case  ,      hypothesis_transform=_snake_case							)['''wer''']
      return {"wer": wer, "cer": cer}
   else:

      def compute_score(_snake_case  :Tuple  ,      _snake_case  :Any  ,      _snake_case  :str="wer"							):
         _A						     =					0
         _A						     =					0
         for prediction, reference in zip(_snake_case  ,      _snake_case							):
            if score_type == "wer":
               _A						     =					compute_measures(_snake_case  ,      _snake_case							)
            elif score_type == "cer":
               _A						     =					compute_measures(
                   _snake_case  ,      _snake_case  ,      truth_transform=_snake_case  ,      hypothesis_transform=_snake_case							)
            incorrect += measures["substitutions"] + measures["deletions"] + measures["insertions"]
            total += measures["substitutions"] + measures["deletions"] + measures["hits"]
         return incorrect / total

      return {"wer": compute_score(_snake_case  ,      _snake_case  ,      '''wer'''							), "cer": compute_score(_snake_case  ,      _snake_case  ,      '''cer'''							)}


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



 """simple docstring"""

 def 						snake_case_			(						self				:				str			)      ->  List[Any]:
    if self.config_name not in _CONFIG_NAMES:
       raise KeyError(f'''You should supply a configuration name selected in {_CONFIG_NAMES}'''			)

    _A						     =					'''int64''' if self.config_name in ['''fleurs-lang_id''', '''minds14'''] else '''string'''

    return datasets.MetricInfo(
        description=_DESCRIPTION ,	citation=_CITATION ,	inputs_description=_KWARGS_DESCRIPTION ,	features=datasets.Features(
            {'''predictions''': datasets.Value(__lowerCAmelCase			), '''references''': datasets.Value(__lowerCAmelCase			)}			) ,	codebase_urls=[] ,	reference_urls=[] ,	format='''numpy''' ,	)

 def 						snake_case_			(						self				:				Dict ,	__lowerCAmelCase				:				str ,	__lowerCAmelCase				:				Any ,	__lowerCAmelCase				:				Union[str, Any]=None ,	__lowerCAmelCase				:				Any=None			)      ->  Union[str, Any]:
    _A						     =					bleu_kwargs if bleu_kwargs is not None else {}
    _A						     =					wer_kwargs if wer_kwargs is not None else {}

    if self.config_name == "fleurs-lang_id":
       return {"accuracy": simple_accuracy(__lowerCAmelCase ,	__lowerCAmelCase			)}
    elif self.config_name == "minds14":
       return fa_and_simple_accuracy(__lowerCAmelCase ,	__lowerCAmelCase			)
    elif self.config_name == "covost2":
       _A						     =					bleu_kwargs.pop('''smooth_method''' ,	'''exp'''			)
       _A						     =					bleu_kwargs.pop('''smooth_value''' ,	__lowerCAmelCase			)
       _A						     =					bleu_kwargs.pop('''force''' ,	__lowerCAmelCase			)
       _A						     =					bleu_kwargs.pop('''lowercase''' ,	__lowerCAmelCase			)
       _A						     =					bleu_kwargs.pop('''tokenize''' ,	__lowerCAmelCase			)
       _A						     =					bleu_kwargs.pop('''use_effective_order''' ,	__lowerCAmelCase			)
       return bleu(
           preds=__lowerCAmelCase ,	labels=__lowerCAmelCase ,	smooth_method=__lowerCAmelCase ,	smooth_value=__lowerCAmelCase ,	force=__lowerCAmelCase ,	lowercase=__lowerCAmelCase ,	tokenize=__lowerCAmelCase ,	use_effective_order=__lowerCAmelCase ,	)
    elif self.config_name in ["fleurs-asr", "mls", "voxpopuli", "babel"]:
       _A						     =					wer_kwargs.pop('''concatenate_texts''' ,	__lowerCAmelCase			)
       return wer_and_cer(__lowerCAmelCase ,	__lowerCAmelCase ,	__lowerCAmelCase ,	self.config_name			)
    else:
       raise KeyError(f'''You should supply a configuration name selected in {_CONFIG_NAMES}'''			)



