



"""simple docstring"""




import gc
import unittest

import numpy as np
import torch

from diffusers import (
    AudioDiffusionPipeline,
    AutoencoderKL,
    DDIMScheduler,
    DDPMScheduler,
    DiffusionPipeline,
    Mel,
    UNetaDConditionModel,
    UNetaDModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()






class      UpperCAmelCase_  (    unittest.TestCase						):



      def 						snake_case_ (					self					:     List[str]						):
           super().tearDown()
           gc.collect()
           torch.cuda.empty_cache()



      @property
      def 						snake_case_ (					self					:     str						):
           torch.manual_seed(0						)
           _UpperCAmelCase      :     Union[str, Any]        =   UNetaDModel(
               sample_size=(3_2, 6_4)					,     in_channels=1					,     out_channels=1					,     layers_per_block=2					,     block_out_channels=(1_2_8, 1_2_8)					,     down_block_types=("AttnDownBlock2D", "DownBlock2D")					,     up_block_types=("UpBlock2D", "AttnUpBlock2D")					,     )
           return model



      @property
      def 						snake_case_ (					self					:     str						):
           torch.manual_seed(0						)
           _UpperCAmelCase      :     Optional[Any]        =   UNetaDConditionModel(
               sample_size=(6_4, 3_2)					,     in_channels=1					,     out_channels=1					,     layers_per_block=2					,     block_out_channels=(1_2_8, 1_2_8)					,     down_block_types=("CrossAttnDownBlock2D", "DownBlock2D")					,     up_block_types=("UpBlock2D", "CrossAttnUpBlock2D")					,     cross_attention_dim=1_0					,     )
           return model



      @property
      def 						snake_case_ (					self					:     Union[str, Any]						):
           torch.manual_seed(0						)
           _UpperCAmelCase      :     Any        =   AutoencoderKL(
               sample_size=(1_2_8, 6_4)					,     in_channels=1					,     out_channels=1					,     latent_channels=1					,     layers_per_block=2					,     block_out_channels=(1_2_8, 1_2_8)					,     down_block_types=("DownEncoderBlock2D", "DownEncoderBlock2D")					,     up_block_types=("UpDecoderBlock2D", "UpDecoderBlock2D")					,     )
           _UpperCAmelCase      :     Union[str, Any]        =   UNetaDModel(
               sample_size=(6_4, 3_2)					,     in_channels=1					,     out_channels=1					,     layers_per_block=2					,     block_out_channels=(1_2_8, 1_2_8)					,     down_block_types=("AttnDownBlock2D", "DownBlock2D")					,     up_block_types=("UpBlock2D", "AttnUpBlock2D")					,     )
           return vqvae, unet



      @slow
      def 						snake_case_ (					self					:     Optional[Any]						):
           _UpperCAmelCase      :     List[Any]        =   '''cpu'''  # ensure determinism for the device-dependent torch.Generator
           _UpperCAmelCase      :     List[Any]        =   Mel(
               x_res=self.dummy_unet.config.sample_size[1]					,     y_res=self.dummy_unet.config.sample_size[0]					,     )

           _UpperCAmelCase      :     Dict        =   DDPMScheduler()
           _UpperCAmelCase      :     List[Any]        =   AudioDiffusionPipeline(vqvae=lowerCamelCase__					,     unet=self.dummy_unet					,     mel=lowerCamelCase__					,     scheduler=lowerCamelCase__						)
           _UpperCAmelCase      :     Optional[int]        =   pipe.to(lowerCamelCase__						)
           pipe.set_progress_bar_config(disable=lowerCamelCase__						)

           _UpperCAmelCase      :     Union[str, Any]        =   torch.Generator(device=lowerCamelCase__						).manual_seed(4_2						)
           _UpperCAmelCase      :     List[str]        =   pipe(generator=lowerCamelCase__					,     steps=4						)
           _UpperCAmelCase      :     Union[str, Any]        =   output.audios[0]
           _UpperCAmelCase      :     Any        =   output.images[0]

           _UpperCAmelCase      :     Optional[int]        =   torch.Generator(device=lowerCamelCase__						).manual_seed(4_2						)
           _UpperCAmelCase      :     int        =   pipe(generator=lowerCamelCase__					,     steps=4					,     return_dict=lowerCamelCase__						)
           _UpperCAmelCase      :     Optional[int]        =   output[0][0]

           assert audio.shape == (1, (self.dummy_unet.config.sample_size[1] - 1) * mel.hop_length)
           assert (
               image.height == self.dummy_unet.config.sample_size[0]
               and image.width == self.dummy_unet.config.sample_size[1]
           )
           _UpperCAmelCase      :     Optional[Any]        =   np.frombuffer(image.tobytes()					,     dtype="uint8"						)[:1_0]
           _UpperCAmelCase      :     Any        =   np.frombuffer(image_from_tuple.tobytes()					,     dtype="uint8"						)[:1_0]
           _UpperCAmelCase      :     str        =   np.array([6_9, 2_5_5, 2_5_5, 2_5_5, 0, 0, 7_7, 1_8_1, 1_2, 1_2_7]						)

           assert np.abs(image_slice.flatten() - expected_slice						).max() == 0
           assert np.abs(image_from_tuple_slice.flatten() - expected_slice						).max() == 0

           _UpperCAmelCase      :     List[str]        =   Mel(
               x_res=self.dummy_vqvae_and_unet[0].config.sample_size[1]					,     y_res=self.dummy_vqvae_and_unet[0].config.sample_size[0]					,     )

           _UpperCAmelCase      :     Optional[Any]        =   DDIMScheduler()
           _UpperCAmelCase      :     Union[str, Any]        =   self.dummy_vqvae_and_unet
           _UpperCAmelCase      :     Union[str, Any]        =   AudioDiffusionPipeline(
               vqvae=self.dummy_vqvae_and_unet[0]					,     unet=dummy_vqvae_and_unet[1]					,     mel=lowerCamelCase__					,     scheduler=lowerCamelCase__						)
           _UpperCAmelCase      :     str        =   pipe.to(lowerCamelCase__						)
           pipe.set_progress_bar_config(disable=lowerCamelCase__						)

           np.random.seed(0						)
           _UpperCAmelCase      :     Optional[int]        =   np.random.uniform(-1					,     1					,     ((dummy_vqvae_and_unet[0].config.sample_size[1] - 1) * mel.hop_length,)						)
           _UpperCAmelCase      :     Optional[int]        =   torch.Generator(device=lowerCamelCase__						).manual_seed(4_2						)
           _UpperCAmelCase      :     Union[str, Any]        =   pipe(raw_audio=lowerCamelCase__					,     generator=lowerCamelCase__					,     start_step=5					,     steps=1_0						)
           _UpperCAmelCase      :     int        =   output.images[0]

           assert (
               image.height == self.dummy_vqvae_and_unet[0].config.sample_size[0]
               and image.width == self.dummy_vqvae_and_unet[0].config.sample_size[1]
           )
           _UpperCAmelCase      :     List[Any]        =   np.frombuffer(image.tobytes()					,     dtype="uint8"						)[:1_0]
           _UpperCAmelCase      :     Union[str, Any]        =   np.array([1_2_0, 1_1_7, 1_1_0, 1_0_9, 1_3_8, 1_6_7, 1_3_8, 1_4_8, 1_3_2, 1_2_1]						)

           assert np.abs(image_slice.flatten() - expected_slice						).max() == 0

           _UpperCAmelCase      :     Any        =   self.dummy_unet_condition
           _UpperCAmelCase      :     str        =   AudioDiffusionPipeline(
               vqvae=self.dummy_vqvae_and_unet[0]					,     unet=lowerCamelCase__					,     mel=lowerCamelCase__					,     scheduler=lowerCamelCase__						)
           _UpperCAmelCase      :     Any        =   pipe.to(lowerCamelCase__						)
           pipe.set_progress_bar_config(disable=lowerCamelCase__						)

           np.random.seed(0						)
           _UpperCAmelCase      :     Union[str, Any]        =   torch.rand((1, 1, 1_0)						)
           _UpperCAmelCase      :     int        =   pipe(generator=lowerCamelCase__					,     encoding=lowerCamelCase__						)
           _UpperCAmelCase      :     Optional[int]        =   output.images[0]
           _UpperCAmelCase      :     List[Any]        =   np.frombuffer(image.tobytes()					,     dtype="uint8"						)[:1_0]
           _UpperCAmelCase      :     str        =   np.array([1_0_7, 1_0_3, 1_2_0, 1_2_7, 1_4_2, 1_2_2, 1_1_3, 1_2_2, 9_7, 1_1_1]						)

           assert np.abs(image_slice.flatten() - expected_slice						).max() == 0







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



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



      def 						snake_case_ (					self					:     str						):
           _UpperCAmelCase      :     Union[str, Any]        =   torch_device

           _UpperCAmelCase      :     Dict        =   DiffusionPipeline.from_pretrained("teticio/audio-diffusion-ddim-256"						)
           _UpperCAmelCase      :     Optional[int]        =   pipe.to(lowerCamelCase__						)
           pipe.set_progress_bar_config(disable=lowerCamelCase__						)

           _UpperCAmelCase      :     Optional[int]        =   torch.Generator(device=lowerCamelCase__						).manual_seed(4_2						)
           _UpperCAmelCase      :     Optional[Any]        =   pipe(generator=lowerCamelCase__						)
           _UpperCAmelCase      :     Dict        =   output.audios[0]
           _UpperCAmelCase      :     Optional[Any]        =   output.images[0]

           assert audio.shape == (1, (pipe.unet.config.sample_size[1] - 1) * pipe.mel.hop_length)
           assert image.height == pipe.unet.config.sample_size[0] and image.width == pipe.unet.config.sample_size[1]
           _UpperCAmelCase      :     Dict        =   np.frombuffer(image.tobytes()					,     dtype="uint8"						)[:1_0]
           _UpperCAmelCase      :     List[Any]        =   np.array([1_5_1, 1_6_7, 1_5_4, 1_4_4, 1_2_2, 1_3_4, 1_2_1, 1_0_5, 7_0, 2_6]						)

           assert np.abs(image_slice.flatten() - expected_slice						).max() == 0









"""simple docstring"""




import io
import json
import unittest

from parameterized import parameterized

from transformers import FSMTForConditionalGeneration, FSMTTokenizer
from transformers.testing_utils import get_tests_dir, require_torch, slow, torch_device
from utils import calculate_bleu


_lowerCAmelCase    : Union[str, Any]							=							get_tests_dir() + "/test_data/fsmt/fsmt_val_data.json"
with io.open(filename, "r", encoding="utf-8") as f:
 _lowerCAmelCase    : Tuple							=							json.load(f)






@require_torch
class      UpperCAmelCase_  (    unittest.TestCase						):



      def 						snake_case_ (					self					:     str					,     A					:     Union[str, Any]						):
           return FSMTTokenizer.from_pretrained(A						)



      def 						snake_case_ (					self					:     Union[str, Any]					,     A					:     Union[str, Any]						):
           _UpperCAmelCase      :     List[Any]        =   FSMTForConditionalGeneration.from_pretrained(A						).to(A						)
           if torch_device == "cuda":
                model.half()
           return model



      @parameterized.expand(
          [
              ["en-ru", 26.0],
              ["ru-en", 22.0],
              ["en-de", 22.0],
              ["de-en", 29.0],
          ]						)
      @slow
      def 						snake_case_ (					self					:     Any					,     A					:     Dict					,     A					:     List[str]						):
           # note: this test is not testing the best performance since it only evals a small batch
           # but it should be enough to detect a regression in the output quality
           _UpperCAmelCase      :     Any        =   f'facebook/wmt19-{pair}'
           _UpperCAmelCase      :     Dict        =   self.get_tokenizer(A						)
           _UpperCAmelCase      :     Optional[int]        =   self.get_model(A						)

           _UpperCAmelCase      :     int        =   bleu_data[pair]["src"]
           _UpperCAmelCase      :     Optional[int]        =   bleu_data[pair]["tgt"]

           _UpperCAmelCase      :     List[str]        =   tokenizer(A					,     return_tensors="pt"					,     truncation=A					,     padding="longest"						).to(A						)
           _UpperCAmelCase      :     List[str]        =   model.generate(
               input_ids=batch.input_ids					,     num_beams=8					,     )
           _UpperCAmelCase      :     Any        =   tokenizer.batch_decode(
               A					,     skip_special_tokens=A					,     clean_up_tokenization_spaces=A						)
           _UpperCAmelCase      :     Any        =   calculate_bleu(A					,     A						)
           print(A						)
           self.assertGreaterEqual(scores["bleu"]					,     A						)
