






import contextlib
import importlib
import io
import unittest

import transformers

# Try to import everything from transformers to ensure every object can be loaded.
from transformers import *  # noqa F406
from transformers.testing_utils import DUMMY_UNKNOWN_IDENTIFIER, require_flax, require_tf, require_torch
from transformers.utils import ContextManagers, find_labels, is_flax_available, is_tf_available, is_torch_available


if is_torch_available():
 from transformers import BertForPreTraining, BertForQuestionAnswering, BertForSequenceClassification

if is_tf_available():
 from transformers import TFBertForPreTraining, TFBertForQuestionAnswering, TFBertForSequenceClassification

if is_flax_available():
 from transformers import FlaxBertForPreTraining, FlaxBertForQuestionAnswering, FlaxBertForSequenceClassification


_a    :   List[str]        =			DUMMY_UNKNOWN_IDENTIFIER
# An actual model hosted on huggingface.co

_a    :   Tuple        =			"""main"""
# Default branch name
_a    :   Dict        =			"""f2c752cfc5c0ab6f4bdec59acea69eefbee381c2"""
# One particular commit (not the top of `main`)
_a    :   Optional[Any]        =			"""aaaaaaa"""
# This commit does not exist, so we should 404.

_a    :   int        =			"""d9e9f15bc825e4b2c9249e9578f884bbcb5e3684"""
# Sha-1 of config.json on the top of `main`, for checking purposes
_a    :   List[str]        =			"""4b243c475af8d0a7754e87d7d096c92e5199ec2fe168a2ee7998e3b8e9bcb1d3"""


@contextlib.contextmanager
def 				snake_case__							(							):
    print("Welcome!"       )
    yield
    print("Bye!"       )


@contextlib.contextmanager
def 				snake_case__							(							):
    print("Bonjour!"       )
    yield
    print("Au revoir!"       )







class 	_UpperCAmelCase  (  unittest.TestCase		):




      """simple docstring"""



      def 			snake_case_      (   self		):



          '''simple docstring'''
          assert transformers.__spec__ is not None
          assert importlib.util.find_spec("transformers"		) is not None







class 	_UpperCAmelCase  (  unittest.TestCase		):




      """simple docstring"""



      @unittest.mock.patch("sys.stdout"    ,	new_callable=io.StringIO		)
      def 			snake_case_      (   self    ,	_lowerCAmelCase		):



          '''simple docstring'''






          with ContextManagers([]		):
              print("Transformers are awesome!"		)
          # The print statement adds a new line at the end of the output
          self.assertEqual(mock_stdout.getvalue()    ,	"Transformers are awesome!\n"		)
      @unittest.mock.patch("sys.stdout"    ,	new_callable=io.StringIO		)
      def 			snake_case_      (   self    ,	_lowerCAmelCase		):



          '''simple docstring'''






          with ContextManagers([context_en()]		):
              print("Transformers are awesome!"		)
          # The output should be wrapped with an English welcome and goodbye
          self.assertEqual(mock_stdout.getvalue()    ,	"Welcome!\nTransformers are awesome!\nBye!\n"		)
      @unittest.mock.patch("sys.stdout"    ,	new_callable=io.StringIO		)
      def 			snake_case_      (   self    ,	_lowerCAmelCase		):



          '''simple docstring'''






          with ContextManagers([context_fr(), context_en()]		):
              print("Transformers are awesome!"		)
          # The output should be wrapped with an English and French welcome and goodbye
          self.assertEqual(mock_stdout.getvalue()    ,	"Bonjour!\nWelcome!\nTransformers are awesome!\nBye!\nAu revoir!\n"		)
      @require_torch
      def 			snake_case_      (   self		):



          '''simple docstring'''
          self.assertEqual(find_labels(__a		)    ,	["labels"]		)
          self.assertEqual(find_labels(__a		)    ,	["labels", "next_sentence_label"]		)
          self.assertEqual(find_labels(__a		)    ,	["start_positions", "end_positions"]		)







          class 	_UpperCAmelCase  (  UpperCamelCase_		):




                """simple docstring"""



                pass





          self.assertEqual(find_labels(__a		)    ,	["labels"]		)
      @require_tf
      def 			snake_case_      (   self		):



          '''simple docstring'''
          self.assertEqual(find_labels(__a		)    ,	["labels"]		)
          self.assertEqual(find_labels(__a		)    ,	["labels", "next_sentence_label"]		)
          self.assertEqual(find_labels(__a		)    ,	["start_positions", "end_positions"]		)







          class 	_UpperCAmelCase  (  UpperCamelCase_		):




                """simple docstring"""



                pass





          self.assertEqual(find_labels(__a		)    ,	["labels"]		)



      @require_flax
      def 			snake_case_      (   self		):



          '''simple docstring'''
          self.assertEqual(find_labels(__a		)    ,	[]		)
          self.assertEqual(find_labels(__a		)    ,	[]		)
          self.assertEqual(find_labels(__a		)    ,	[]		)







          class 	_UpperCAmelCase  (  UpperCamelCase_		):




                """simple docstring"""



                pass





          self.assertEqual(find_labels(__a		)    ,	[]		)








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		):




      """simple docstring"""



      A          =    ['''vqvae''']
      def __init__(   self    ,	_lowerCAmelCase    ,	_lowerCAmelCase    ,	_lowerCAmelCase    ,	_lowerCAmelCase    ,	):



          '''simple docstring'''
          super().__init__()
          self.register_modules(unet=_lowerCAmelCase    ,	scheduler=_lowerCAmelCase    ,	mel=_lowerCAmelCase    ,	vqvae=_lowerCAmelCase		)
      def 			snake_case_      (   self		):



          '''simple docstring'''
          return 50 if isinstance(self.scheduler    ,	_lowerCAmelCase		) else 1_000
      @torch.no_grad()
      def __call__(   self    ,	_lowerCAmelCase = 1    ,	_lowerCAmelCase = None    ,	_lowerCAmelCase = None    ,	_lowerCAmelCase = 0    ,	_lowerCAmelCase = 0    ,	_lowerCAmelCase = None    ,	_lowerCAmelCase = None    ,	_lowerCAmelCase = 0    ,	_lowerCAmelCase = 0    ,	_lowerCAmelCase = None    ,	_lowerCAmelCase = 0    ,	_lowerCAmelCase = None    ,	_lowerCAmelCase = None    ,	_lowerCAmelCase=True    ,	):



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

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

              if self.vqvae is not None:
                  lowerCAmelCase__   :str      							=				self.vqvae.encode(torch.unsqueeze(_lowerCAmelCase    ,	0		)		).latent_dist.sample(
                      generator=_lowerCAmelCase		)[0]
                  lowerCAmelCase__   :Dict      							=				self.vqvae.config.scaling_factor * input_images

              if start_step > 0:
                  lowerCAmelCase__   :Dict      							=				self.scheduler.add_noise(_lowerCAmelCase    ,	_lowerCAmelCase    ,	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__   :Dict      							=				int(mask_start_secs * pixels_per_second		)
              lowerCAmelCase__   :Tuple      							=				int(mask_end_secs * pixels_per_second		)
              lowerCAmelCase__   :str      							=				self.scheduler.add_noise(_lowerCAmelCase    ,	_lowerCAmelCase    ,	torch.tensor(self.scheduler.timesteps[start_step:]		)		)

          for step, t in enumerate(self.progress_bar(self.scheduler.timesteps[start_step:]		)		):
              if isinstance(self.unet    ,	_lowerCAmelCase		):
                  lowerCAmelCase__   :Optional[Any]      							=				self.unet(_lowerCAmelCase    ,	_lowerCAmelCase    ,	_lowerCAmelCase		)["sample"]
              else:
                  lowerCAmelCase__   :Dict      							=				self.unet(_lowerCAmelCase    ,	_lowerCAmelCase		)["sample"]

              if isinstance(self.scheduler    ,	_lowerCAmelCase		):
                  lowerCAmelCase__   :Any      							=				self.scheduler.step(
                      model_output=_lowerCAmelCase    ,	timestep=_lowerCAmelCase    ,	sample=_lowerCAmelCase    ,	eta=_lowerCAmelCase    ,	generator=_lowerCAmelCase    ,	)["prev_sample"]
              else:
                  lowerCAmelCase__   :List[str]      							=				self.scheduler.step(
                      model_output=_lowerCAmelCase    ,	timestep=_lowerCAmelCase    ,	sample=_lowerCAmelCase    ,	generator=_lowerCAmelCase    ,	)["prev_sample"]

              if mask is not None:
                  if mask_start > 0:
                      lowerCAmelCase__   :List[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__   :Any      							=				1 / self.vqvae.config.scaling_factor * images
              lowerCAmelCase__   :List[Any]      							=				self.vqvae.decode(_lowerCAmelCase		)["sample"]

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

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

          return BaseOutput(**AudioPipelineOutput(np.array(_lowerCAmelCase		)[:, np.newaxis, :]		)    ,	**ImagePipelineOutput(_lowerCAmelCase		)		)
      @torch.no_grad()
      def 			snake_case_      (   self    ,	_lowerCAmelCase    ,	_lowerCAmelCase = 50		):



          '''simple docstring'''
          assert isinstance(self.scheduler    ,	_lowerCAmelCase		)
          self.scheduler.set_timesteps(_lowerCAmelCase		)
          lowerCAmelCase__   :Any      							=				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__   :Optional[Any]      							=				torch.Tensor(_lowerCAmelCase		).to(self.device		)

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

          return sample



      @staticmethod
      def 			snake_case_      (   _lowerCAmelCase    ,	_lowerCAmelCase    ,	_lowerCAmelCase		):



          '''simple docstring'''
          lowerCAmelCase__   :List[str]      							=				acos(torch.dot(torch.flatten(_lowerCAmelCase		)    ,	torch.flatten(_lowerCAmelCase		)		) / torch.norm(_lowerCAmelCase		) / torch.norm(_lowerCAmelCase		)		)
          return sin((1 - alpha) * theta		) * xa / sin(_lowerCAmelCase		) + sin(alpha * theta		) * xa / sin(_lowerCAmelCase		)



