

"""simple docstring"""






import gc
import unittest

from transformers import MODEL_FOR_MASKED_LM_MAPPING, TF_MODEL_FOR_MASKED_LM_MAPPING, FillMaskPipeline, pipeline
from transformers.pipelines import PipelineException
from transformers.testing_utils import (
    is_pipeline_test,
    is_torch_available,
    nested_simplify,
    require_tf,
    require_torch,
    require_torch_gpu,
    slow,
)

from .test_pipelines_common import ANY







@is_pipeline_test
class 	__UpperCamelCase    (    unittest.TestCase						):
  _UpperCAmelCase          =		MODEL_FOR_MASKED_LM_MAPPING
  _UpperCAmelCase          =		TF_MODEL_FOR_MASKED_LM_MAPPING

  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      super().tearDown()
      # clean-up as much as possible GPU memory occupied by PyTorch
      gc.collect()
      if is_torch_available():
          import torch

          torch.cuda.empty_cache()

  @require_tf
  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase       :   int  =    pipeline(task='fill-mask'    ,model='sshleifer/tiny-distilroberta-base'    ,top_k=2    ,framework='tf'      )
      _lowerCAmelCase       :   Any  =    unmasker('My name is <mask>'      )
      self.assertEqual(
          nested_simplify(_A    ,decimals=6      )    ,[
              {'sequence': 'My name is grouped', 'score': 2.1E-05, 'token': 3_8015, 'token_str': ' grouped'},
              {'sequence': 'My name is accuser', 'score': 2.1E-05, 'token': 2_5506, 'token_str': ' accuser'},
          ]    ,)

      _lowerCAmelCase       :   List[Any]  =    unmasker('The largest city in France is <mask>'      )
      self.assertEqual(
          nested_simplify(_A    ,decimals=6      )    ,[
              {
                  'sequence': 'The largest city in France is grouped',
                  'score': 2.1E-05,
                  'token': 3_8015,
                  'token_str': ' grouped',
              },
              {
                  'sequence': 'The largest city in France is accuser',
                  'score': 2.1E-05,
                  'token': 2_5506,
                  'token_str': ' accuser',
              },
          ]    ,)

      _lowerCAmelCase       :   Optional[int]  =    unmasker('My name is <mask>'    ,targets=[' Patrick', ' Clara', ' Teven']    ,top_k=3      )
      self.assertEqual(
          nested_simplify(_A    ,decimals=6      )    ,[
              {'sequence': 'My name is Clara', 'score': 2E-05, 'token': 1_3606, 'token_str': ' Clara'},
              {'sequence': 'My name is Patrick', 'score': 2E-05, 'token': 3499, 'token_str': ' Patrick'},
              {'sequence': 'My name is Te', 'score': 1.9E-05, 'token': 2941, 'token_str': ' Te'},
          ]    ,)

  @require_torch
  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase       :   List[str]  =    pipeline(task='fill-mask'    ,model='sshleifer/tiny-distilroberta-base'    ,top_k=2    ,framework='pt'      )

      _lowerCAmelCase       :   int  =    unmasker('My name is <mask>'      )
      self.assertEqual(
          nested_simplify(_A    ,decimals=6      )    ,[
              {'sequence': 'My name is Maul', 'score': 2.2E-05, 'token': 3_5676, 'token_str': ' Maul'},
              {'sequence': 'My name isELS', 'score': 2.2E-05, 'token': 1_6416, 'token_str': 'ELS'},
          ]    ,)

      _lowerCAmelCase       :   Union[str, Any]  =    unmasker('The largest city in France is <mask>'      )
      self.assertEqual(
          nested_simplify(_A    ,decimals=6      )    ,[
              {
                  'sequence': 'The largest city in France is Maul',
                  'score': 2.2E-05,
                  'token': 3_5676,
                  'token_str': ' Maul',
              },
              {'sequence': 'The largest city in France isELS', 'score': 2.2E-05, 'token': 1_6416, 'token_str': 'ELS'},
          ]    ,)

      _lowerCAmelCase       :   Tuple  =    unmasker('My name is <mask>'    ,targets=[' Patrick', ' Clara', ' Teven']    ,top_k=3      )
      self.assertEqual(
          nested_simplify(_A    ,decimals=6      )    ,[
              {'sequence': 'My name is Patrick', 'score': 2.1E-05, 'token': 3499, 'token_str': ' Patrick'},
              {'sequence': 'My name is Te', 'score': 2E-05, 'token': 2941, 'token_str': ' Te'},
              {'sequence': 'My name is Clara', 'score': 2E-05, 'token': 1_3606, 'token_str': ' Clara'},
          ]    ,)

      _lowerCAmelCase       :   int  =    unmasker('My name is <mask> <mask>'    ,top_k=2      )

      self.assertEqual(
          nested_simplify(_A    ,decimals=6      )    ,[
              [
                  {
                      'score': 2.2E-05,
                      'token': 3_5676,
                      'token_str': ' Maul',
                      'sequence': '<s>My name is Maul<mask></s>',
                  },
                  {'score': 2.2E-05, 'token': 1_6416, 'token_str': 'ELS', 'sequence': '<s>My name isELS<mask></s>'},
              ],
              [
                  {
                      'score': 2.2E-05,
                      'token': 3_5676,
                      'token_str': ' Maul',
                      'sequence': '<s>My name is<mask> Maul</s>',
                  },
                  {'score': 2.2E-05, 'token': 1_6416, 'token_str': 'ELS', 'sequence': '<s>My name is<mask>ELS</s>'},
              ],
          ]    ,)

  @require_torch_gpu
  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase       :   Tuple  =    pipeline('fill-mask'    ,model='hf-internal-testing/tiny-random-distilbert'    ,device=0    ,framework='pt'      )

      # convert model to fp16
      pipe.model.half()

      _lowerCAmelCase       :   Tuple  =    pipe('Paris is the [MASK] of France.'      )
      # We actually don't care about the result, we just want to make sure
      # it works, meaning the float16 tensor got casted back to float32
      # for postprocessing.
      self.assertIsInstance(_A    ,_A      )

  @slow
  @require_torch
  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase       :   Optional[Any]  =    pipeline(task='fill-mask'    ,model='distilroberta-base'    ,top_k=2    ,framework='pt'      )
      self.run_large_test(_A      )

  @slow
  @require_tf
  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase       :   Union[str, Any]  =    pipeline(task='fill-mask'    ,model='distilroberta-base'    ,top_k=2    ,framework='tf'      )
      self.run_large_test(_A      )

  def __lowerCamelCase		(		self    ,_A      ):
      '''simple docstring'''

      _lowerCAmelCase       :   int  =    unmasker('My name is <mask>'      )
      self.assertEqual(
          nested_simplify(_A      )    ,[
              {'sequence': 'My name is John', 'score': 0.0_0_8, 'token': 610, 'token_str': ' John'},
              {'sequence': 'My name is Chris', 'score': 0.0_0_7, 'token': 1573, 'token_str': ' Chris'},
          ]    ,)
      _lowerCAmelCase       :   List[Any]  =    unmasker('The largest city in France is <mask>'      )
      self.assertEqual(
          nested_simplify(_A      )    ,[
              {
                  'sequence': 'The largest city in France is Paris',
                  'score': 0.2_5_1,
                  'token': 2201,
                  'token_str': ' Paris',
              },
              {
                  'sequence': 'The largest city in France is Lyon',
                  'score': 0.2_1_4,
                  'token': 1_2790,
                  'token_str': ' Lyon',
              },
          ]    ,)

      _lowerCAmelCase       :   Optional[int]  =    unmasker('My name is <mask>'    ,targets=[' Patrick', ' Clara', ' Teven']    ,top_k=3      )
      self.assertEqual(
          nested_simplify(_A      )    ,[
              {'sequence': 'My name is Patrick', 'score': 0.0_0_5, 'token': 3499, 'token_str': ' Patrick'},
              {'sequence': 'My name is Clara', 'score': 0.0_0_0, 'token': 1_3606, 'token_str': ' Clara'},
              {'sequence': 'My name is Te', 'score': 0.0_0_0, 'token': 2941, 'token_str': ' Te'},
          ]    ,)

  @require_torch
  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase       :   Dict  =    pipeline(task='fill-mask'    ,model='sshleifer/tiny-distilroberta-base'    ,framework='pt'      )
      _lowerCAmelCase       :   Optional[int]  =    None
      _lowerCAmelCase       :   Optional[int]  =    None
      self.run_pipeline_test(_A    ,[]      )

  @require_tf
  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase       :   Union[str, Any]  =    pipeline(task='fill-mask'    ,model='sshleifer/tiny-distilroberta-base'    ,framework='tf'      )
      _lowerCAmelCase       :   List[Any]  =    None
      _lowerCAmelCase       :   Dict  =    None
      self.run_pipeline_test(_A    ,[]      )

  def __lowerCamelCase		(		self    ,_A    ,_A    ,_A      ):
      '''simple docstring'''

      if tokenizer is None or tokenizer.mask_token_id is None:
          self.skipTest('The provided tokenizer has no mask token, (probably reformer or wav2vec2)'      )

      _lowerCAmelCase       :   List[str]  =    FillMaskPipeline(model=_A    ,tokenizer=_A      )
      _lowerCAmelCase       :   List[str]  =    [
          F"""This is another {tokenizer.mask_token} test""",
      ]
      return fill_masker, examples

  def __lowerCamelCase		(		self    ,_A    ,_A      ):
      '''simple docstring'''

      _lowerCAmelCase       :   Tuple  =    fill_masker.tokenizer
      _lowerCAmelCase       :   Optional[int]  =    fill_masker.model

      _lowerCAmelCase       :   Union[str, Any]  =    fill_masker(
          F"""This is a {tokenizer.mask_token}"""    ,)
      self.assertEqual(
          _A    ,[
              {'sequence': ANY(_A      ), 'score': ANY(_A      ), 'token': ANY(_A      ), 'token_str': ANY(_A      )},
              {'sequence': ANY(_A      ), 'score': ANY(_A      ), 'token': ANY(_A      ), 'token_str': ANY(_A      )},
              {'sequence': ANY(_A      ), 'score': ANY(_A      ), 'token': ANY(_A      ), 'token_str': ANY(_A      )},
              {'sequence': ANY(_A      ), 'score': ANY(_A      ), 'token': ANY(_A      ), 'token_str': ANY(_A      )},
              {'sequence': ANY(_A      ), 'score': ANY(_A      ), 'token': ANY(_A      ), 'token_str': ANY(_A      )},
          ]    ,)

      _lowerCAmelCase       :   Any  =    fill_masker([F"""This is a {tokenizer.mask_token}"""]      )
      self.assertEqual(
          _A    ,[
              {'sequence': ANY(_A      ), 'score': ANY(_A      ), 'token': ANY(_A      ), 'token_str': ANY(_A      )},
              {'sequence': ANY(_A      ), 'score': ANY(_A      ), 'token': ANY(_A      ), 'token_str': ANY(_A      )},
              {'sequence': ANY(_A      ), 'score': ANY(_A      ), 'token': ANY(_A      ), 'token_str': ANY(_A      )},
              {'sequence': ANY(_A      ), 'score': ANY(_A      ), 'token': ANY(_A      ), 'token_str': ANY(_A      )},
              {'sequence': ANY(_A      ), 'score': ANY(_A      ), 'token': ANY(_A      ), 'token_str': ANY(_A      )},
          ]    ,)

      _lowerCAmelCase       :   Tuple  =    fill_masker([F"""This is a {tokenizer.mask_token}""", F"""Another {tokenizer.mask_token} great test."""]      )
      self.assertEqual(
          _A    ,[
              [
                  {'sequence': ANY(_A      ), 'score': ANY(_A      ), 'token': ANY(_A      ), 'token_str': ANY(_A      )},
                  {'sequence': ANY(_A      ), 'score': ANY(_A      ), 'token': ANY(_A      ), 'token_str': ANY(_A      )},
                  {'sequence': ANY(_A      ), 'score': ANY(_A      ), 'token': ANY(_A      ), 'token_str': ANY(_A      )},
                  {'sequence': ANY(_A      ), 'score': ANY(_A      ), 'token': ANY(_A      ), 'token_str': ANY(_A      )},
                  {'sequence': ANY(_A      ), 'score': ANY(_A      ), 'token': ANY(_A      ), 'token_str': ANY(_A      )},
              ],
              [
                  {'sequence': ANY(_A      ), 'score': ANY(_A      ), 'token': ANY(_A      ), 'token_str': ANY(_A      )},
                  {'sequence': ANY(_A      ), 'score': ANY(_A      ), 'token': ANY(_A      ), 'token_str': ANY(_A      )},
                  {'sequence': ANY(_A      ), 'score': ANY(_A      ), 'token': ANY(_A      ), 'token_str': ANY(_A      )},
                  {'sequence': ANY(_A      ), 'score': ANY(_A      ), 'token': ANY(_A      ), 'token_str': ANY(_A      )},
                  {'sequence': ANY(_A      ), 'score': ANY(_A      ), 'token': ANY(_A      ), 'token_str': ANY(_A      )},
              ],
          ]    ,)

      with self.assertRaises(_A      ):
          fill_masker([None]      )
      # No mask_token is not supported
      with self.assertRaises(_A      ):
          fill_masker('This is'      )

      self.run_test_top_k(_A    ,_A      )
      self.run_test_targets(_A    ,_A      )
      self.run_test_top_k_targets(_A    ,_A      )
      self.fill_mask_with_duplicate_targets_and_top_k(_A    ,_A      )
      self.fill_mask_with_multiple_masks(_A    ,_A      )

  def __lowerCamelCase		(		self    ,_A    ,_A      ):
      '''simple docstring'''

      _lowerCAmelCase       :   Optional[Any]  =    tokenizer.get_vocab()
      _lowerCAmelCase       :   Tuple  =    sorted(vocab.keys()      )[:2]
      # Pipeline argument
      _lowerCAmelCase       :   Optional[Any]  =    FillMaskPipeline(model=_A    ,tokenizer=_A    ,targets=_A      )
      _lowerCAmelCase       :   Tuple  =    fill_masker(F"""This is a {tokenizer.mask_token}"""      )
      self.assertEqual(
          _A    ,[
              {'sequence': ANY(_A      ), 'score': ANY(_A      ), 'token': ANY(_A      ), 'token_str': ANY(_A      )},
              {'sequence': ANY(_A      ), 'score': ANY(_A      ), 'token': ANY(_A      ), 'token_str': ANY(_A      )},
          ]    ,)
      _lowerCAmelCase       :   Union[str, Any]  =    {vocab[el] for el in targets}
      self.assertEqual({el['token'] for el in outputs}    ,_A      )
      _lowerCAmelCase       :   List[Any]  =    [tokenizer.decode([x]      ) for x in target_ids]
      self.assertEqual({el['token_str'] for el in outputs}    ,set(_A      )      )

      # Call argument
      _lowerCAmelCase       :   Optional[Any]  =    FillMaskPipeline(model=_A    ,tokenizer=_A      )
      _lowerCAmelCase       :   str  =    fill_masker(F"""This is a {tokenizer.mask_token}"""    ,targets=_A      )
      self.assertEqual(
          _A    ,[
              {'sequence': ANY(_A      ), 'score': ANY(_A      ), 'token': ANY(_A      ), 'token_str': ANY(_A      )},
              {'sequence': ANY(_A      ), 'score': ANY(_A      ), 'token': ANY(_A      ), 'token_str': ANY(_A      )},
          ]    ,)
      _lowerCAmelCase       :   int  =    {vocab[el] for el in targets}
      self.assertEqual({el['token'] for el in outputs}    ,_A      )
      _lowerCAmelCase       :   Dict  =    [tokenizer.decode([x]      ) for x in target_ids]
      self.assertEqual({el['token_str'] for el in outputs}    ,set(_A      )      )

      # Score equivalence
      _lowerCAmelCase       :   Any  =    fill_masker(F"""This is a {tokenizer.mask_token}"""    ,targets=_A      )
      _lowerCAmelCase       :   Any  =    [top_mask['token_str'] for top_mask in outputs]
      _lowerCAmelCase       :   Optional[int]  =    [top_mask['score'] for top_mask in outputs]

      # For some BPE tokenizers, `</w>` is removed during decoding, so `token_str` won't be the same as in `targets`.
      if set(_A      ) == set(_A      ):
          _lowerCAmelCase       :   Dict  =    fill_masker(F"""This is a {tokenizer.mask_token}"""    ,targets=_A      )
          _lowerCAmelCase       :   Union[str, Any]  =    [top_mask['score'] for top_mask in unmasked_targets]
          self.assertEqual(nested_simplify(_A      )    ,nested_simplify(_A      )      )

      # Raises with invalid
      with self.assertRaises(_A      ):
          _lowerCAmelCase       :   Any  =    fill_masker(F"""This is a {tokenizer.mask_token}"""    ,targets=[]      )
      # For some tokenizers, `""` is actually in the vocabulary and the expected error won't raised
      if "" not in tokenizer.get_vocab():
          with self.assertRaises(_A      ):
              _lowerCAmelCase       :   int  =    fill_masker(F"""This is a {tokenizer.mask_token}"""    ,targets=['']      )
          with self.assertRaises(_A      ):
              _lowerCAmelCase       :   Dict  =    fill_masker(F"""This is a {tokenizer.mask_token}"""    ,targets=''      )

  def __lowerCamelCase		(		self    ,_A    ,_A      ):
      '''simple docstring'''

      _lowerCAmelCase       :   Optional[int]  =    FillMaskPipeline(model=_A    ,tokenizer=_A    ,top_k=2      )
      _lowerCAmelCase       :   Optional[Any]  =    fill_masker(F"""This is a {tokenizer.mask_token}"""      )
      self.assertEqual(
          _A    ,[
              {'sequence': ANY(_A      ), 'score': ANY(_A      ), 'token': ANY(_A      ), 'token_str': ANY(_A      )},
              {'sequence': ANY(_A      ), 'score': ANY(_A      ), 'token': ANY(_A      ), 'token_str': ANY(_A      )},
          ]    ,)

      _lowerCAmelCase       :   Optional[Any]  =    FillMaskPipeline(model=_A    ,tokenizer=_A      )
      _lowerCAmelCase       :   List[Any]  =    fill_masker(F"""This is a {tokenizer.mask_token}"""    ,top_k=2      )
      self.assertEqual(
          _A    ,[
              {'sequence': ANY(_A      ), 'score': ANY(_A      ), 'token': ANY(_A      ), 'token_str': ANY(_A      )},
              {'sequence': ANY(_A      ), 'score': ANY(_A      ), 'token': ANY(_A      ), 'token_str': ANY(_A      )},
          ]    ,)
      self.assertEqual(nested_simplify(_A      )    ,nested_simplify(_A      )      )

  def __lowerCamelCase		(		self    ,_A    ,_A      ):
      '''simple docstring'''

      _lowerCAmelCase       :   Optional[Any]  =    tokenizer.get_vocab()
      _lowerCAmelCase       :   List[str]  =    FillMaskPipeline(model=_A    ,tokenizer=_A      )

      # top_k=2, ntargets=3
      _lowerCAmelCase       :   str  =    sorted(vocab.keys()      )[:3]
      _lowerCAmelCase       :   Dict  =    fill_masker(F"""This is a {tokenizer.mask_token}"""    ,top_k=2    ,targets=_A      )

      # If we use the most probably targets, and filter differently, we should still
      # have the same results
      _lowerCAmelCase       :   Dict  =    [el['token_str'] for el in sorted(_A    ,key=lambda _A      : x["score"]    ,reverse=_A      )]
      # For some BPE tokenizers, `</w>` is removed during decoding, so `token_str` won't be the same as in `targets`.
      if set(_A      ).issubset(_A      ):
          _lowerCAmelCase       :   List[Any]  =    fill_masker(F"""This is a {tokenizer.mask_token}"""    ,top_k=3    ,targets=_A      )
          # They should yield exactly the same result
          self.assertEqual(nested_simplify(_A      )    ,nested_simplify(_A      )      )

  def __lowerCamelCase		(		self    ,_A    ,_A      ):
      '''simple docstring'''

      _lowerCAmelCase       :   List[Any]  =    FillMaskPipeline(model=_A    ,tokenizer=_A      )
      _lowerCAmelCase       :   Tuple  =    tokenizer.get_vocab()
      # String duplicates + id duplicates
      _lowerCAmelCase       :   List[Any]  =    sorted(vocab.keys()      )[:3]
      _lowerCAmelCase       :   List[Any]  =    [targets[0], targets[1], targets[0], targets[2], targets[1]]
      _lowerCAmelCase       :   List[Any]  =    fill_masker(F"""My name is {tokenizer.mask_token}"""    ,targets=_A    ,top_k=10      )

      # The target list contains duplicates, so we can't output more
      # than them
      self.assertEqual(len(_A      )    ,3      )







  def __lowerCamelCase		(		self    ,_A    ,_A      ):
      '''simple docstring'''

      _lowerCAmelCase       :   List[str]  =    FillMaskPipeline(model=_A    ,tokenizer=_A      )

      _lowerCAmelCase       :   Union[str, Any]  =    fill_masker(
          F"""This is a {tokenizer.mask_token} {tokenizer.mask_token} {tokenizer.mask_token}"""    ,top_k=2      )
      self.assertEqual(
          _A    ,[
              [
                  {'sequence': ANY(_A      ), 'score': ANY(_A      ), 'token': ANY(_A      ), 'token_str': ANY(_A      )},
                  {'sequence': ANY(_A      ), 'score': ANY(_A      ), 'token': ANY(_A      ), 'token_str': ANY(_A      )},
              ],
              [
                  {'sequence': ANY(_A      ), 'score': ANY(_A      ), 'token': ANY(_A      ), 'token_str': ANY(_A      )},
                  {'sequence': ANY(_A      ), 'score': ANY(_A      ), 'token': ANY(_A      ), 'token_str': ANY(_A      )},
              ],
              [
                  {'sequence': ANY(_A      ), 'score': ANY(_A      ), 'token': ANY(_A      ), 'token_str': ANY(_A      )},
                  {'sequence': ANY(_A      ), 'score': ANY(_A      ), 'token': ANY(_A      ), 'token_str': ANY(_A      )},
              ],
          ]    ,)







"""simple docstring"""






from math import acos, sin
from typing import List, Tuple, Union

import numpy as np
import torch
from PIL import Image

from ...models import AutoencoderKL, UNetaDConditionModel
from ...schedulers import DDIMScheduler, DDPMScheduler
from ...utils import randn_tensor
from ..pipeline_utils import AudioPipelineOutput, BaseOutput, DiffusionPipeline, ImagePipelineOutput
from .mel import Mel







class 	__UpperCamelCase    (    a__						):
  _UpperCAmelCase          =		["vqvae"]

  def __init__(		self    ,_A    ,_A    ,_A    ,_A    ,):
      '''simple docstring'''

      super().__init__()
      self.register_modules(unet=_A    ,scheduler=_A    ,mel=_A    ,vqvae=_A      )

  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      return 50 if isinstance(self.scheduler    ,_A      ) else 1000

  @torch.no_grad()
  def __call__(		self    ,_A = 1    ,_A = None    ,_A = None    ,_A = 0    ,_A = 0    ,_A = None    ,_A = None    ,_A = 0    ,_A = 0    ,_A = None    ,_A = 0    ,_A = None    ,_A = None    ,_A=True    ,):
      '''simple docstring'''

      _lowerCAmelCase       :   List[str]  =    steps or self.get_default_steps()
      self.scheduler.set_timesteps(_A      )
      _lowerCAmelCase       :   Optional[Any]  =    step_generator or generator
      # For backwards compatibility
      if type(self.unet.config.sample_size      ) == int:
          _lowerCAmelCase       :   Tuple  =    (self.unet.config.sample_size, self.unet.config.sample_size)
      if noise is None:
          _lowerCAmelCase       :   Optional[Any]  =    randn_tensor(
              (
                  batch_size,
                  self.unet.config.in_channels,
                  self.unet.config.sample_size[0],
                  self.unet.config.sample_size[1],
              )    ,generator=_A    ,device=self.device    ,)
      _lowerCAmelCase       :   Dict  =    noise
      _lowerCAmelCase       :   Optional[Any]  =    None

      if audio_file is not None or raw_audio is not None:
          self.mel.load_audio(_A    ,_A      )
          _lowerCAmelCase       :   Union[str, Any]  =    self.mel.audio_slice_to_image(_A      )
          _lowerCAmelCase       :   int  =    np.frombuffer(input_image.tobytes()    ,dtype='uint8'      ).reshape(
              (input_image.height, input_image.width)      )
          _lowerCAmelCase       :   int  =    (input_image / 255) * 2 - 1
          _lowerCAmelCase       :   str  =    torch.tensor(input_image[np.newaxis, :, :]    ,dtype=torch.float      ).to(self.device      )

          if self.vqvae is not None:
              _lowerCAmelCase       :   List[Any]  =    self.vqvae.encode(torch.unsqueeze(_A    ,0      )      ).latent_dist.sample(
                  generator=_A      )[0]
              _lowerCAmelCase       :   Tuple  =    self.vqvae.config.scaling_factor * input_images

          if start_step > 0:
              _lowerCAmelCase       :   List[Any]  =    self.scheduler.add_noise(_A    ,_A    ,self.scheduler.timesteps[start_step - 1]      )

          _lowerCAmelCase       :   Optional[Any]  =    (
              self.unet.config.sample_size[1] * self.mel.get_sample_rate() / self.mel.x_res / self.mel.hop_length
          )
          _lowerCAmelCase       :   Optional[Any]  =    int(mask_start_secs * pixels_per_second      )
          _lowerCAmelCase       :   Optional[int]  =    int(mask_end_secs * pixels_per_second      )
          _lowerCAmelCase       :   int  =    self.scheduler.add_noise(_A    ,_A    ,torch.tensor(self.scheduler.timesteps[start_step:]      )      )

      for step, t in enumerate(self.progress_bar(self.scheduler.timesteps[start_step:]      )      ):
          if isinstance(self.unet    ,_A      ):
              _lowerCAmelCase       :   str  =    self.unet(_A    ,_A    ,_A      )['sample']
          else:
              _lowerCAmelCase       :   Any  =    self.unet(_A    ,_A      )['sample']

          if isinstance(self.scheduler    ,_A      ):
              _lowerCAmelCase       :   Union[str, Any]  =    self.scheduler.step(
                  model_output=_A    ,timestep=_A    ,sample=_A    ,eta=_A    ,generator=_A    ,)['prev_sample']
          else:
              _lowerCAmelCase       :   Any  =    self.scheduler.step(
                  model_output=_A    ,timestep=_A    ,sample=_A    ,generator=_A    ,)['prev_sample']

          if mask is not None:
              if mask_start > 0:
                  _lowerCAmelCase       :   Any  =    mask[:, step, :, :mask_start]
              if mask_end > 0:
                  _lowerCAmelCase       :   Optional[Any]  =    mask[:, step, :, -mask_end:]

      if self.vqvae is not None:
          # 0.18215 was scaling factor used in training to ensure unit variance
          _lowerCAmelCase       :   Union[str, Any]  =    1 / self.vqvae.config.scaling_factor * images
          _lowerCAmelCase       :   Any  =    self.vqvae.decode(_A      )['sample']

      _lowerCAmelCase       :   Any  =    (images / 2 + 0.5).clamp(0    ,1      )
      _lowerCAmelCase       :   Tuple  =    images.cpu().permute(0    ,2    ,3    ,1      ).numpy()
      _lowerCAmelCase       :   Any  =    (images * 255).round().astype('uint8'      )
      _lowerCAmelCase       :   Any  =    list(
          (Image.fromarray(_[:, :, 0]      ) for _ in images)
          if images.shape[3] == 1
          else (Image.fromarray(_A    ,mode='RGB'      ).convert('L'      ) for _ in images)      )

      _lowerCAmelCase       :   Dict  =    [self.mel.image_to_audio(_A      ) for _ in images]
      if not return_dict:
          return images, (self.mel.get_sample_rate(), audios)

      return BaseOutput(**AudioPipelineOutput(np.array(_A      )[:, np.newaxis, :]      )    ,**ImagePipelineOutput(_A      )      )

  @torch.no_grad()
  def __lowerCamelCase		(		self    ,_A    ,_A = 50      ):
      '''simple docstring'''

      assert isinstance(self.scheduler    ,_A      )
      self.scheduler.set_timesteps(_A      )
      _lowerCAmelCase       :   Dict  =    np.array(
          [np.frombuffer(image.tobytes()    ,dtype='uint8'      ).reshape((1, image.height, image.width)      ) for image in images]      )
      _lowerCAmelCase       :   Dict  =    (sample / 255) * 2 - 1
      _lowerCAmelCase       :   List[str]  =    torch.Tensor(_A      ).to(self.device      )

      for t in self.progress_bar(torch.flip(self.scheduler.timesteps    ,(0,)      )      ):
          _lowerCAmelCase       :   Tuple  =    t - self.scheduler.config.num_train_timesteps // self.scheduler.num_inference_steps
          _lowerCAmelCase       :   Optional[int]  =    self.scheduler.alphas_cumprod[t]
          _lowerCAmelCase       :   Dict  =    (
              self.scheduler.alphas_cumprod[prev_timestep]
              if prev_timestep >= 0
              else self.scheduler.final_alpha_cumprod
          )
          _lowerCAmelCase       :   Union[str, Any]  =    1 - alpha_prod_t
          _lowerCAmelCase       :   Union[str, Any]  =    self.unet(_A    ,_A      )['sample']
          _lowerCAmelCase       :   Optional[int]  =    (1 - alpha_prod_t_prev) ** 0.5 * model_output
          _lowerCAmelCase       :   Any  =    (sample - pred_sample_direction) * alpha_prod_t_prev ** (-0.5)
          _lowerCAmelCase       :   Dict  =    sample * alpha_prod_t ** 0.5 + beta_prod_t ** 0.5 * model_output

      return sample







  @staticmethod
  def __lowerCamelCase		(		_A    ,_A    ,_A      ):
      '''simple docstring'''

      _lowerCAmelCase       :   int  =    acos(torch.dot(torch.flatten(_A      )    ,torch.flatten(_A      )      ) / torch.norm(_A      ) / torch.norm(_A      )      )
      return sin((1 - alpha) * theta      ) * xa / sin(_A      ) + sin(alpha * theta      ) * xa / sin(_A      )





