"""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 					lowerCAmelCase_      (     lowerCAmelCase   ):



				"""simple docstring"""
				_lowerCAmelCase       :      Dict     					=  ["""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 10_00

				@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"""

										snake_case        =     steps or self.get_default_steps()
										self.scheduler.set_timesteps(lowerCAmelCase      )
										snake_case        =     step_generator or generator
										# For backwards compatibility
										if type(self.unet.config.sample_size      ) == int:
																snake_case        =     (self.unet.config.sample_size, self.unet.config.sample_size)
										if noise is None:
																snake_case        =     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						,      )
										snake_case        =     noise
										snake_case        =     None

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

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

																if start_step > 0:
																						snake_case        =     self.scheduler.add_noise(lowerCAmelCase						,      lowerCAmelCase						,      self.scheduler.timesteps[start_step - 1]      )

																snake_case        =     (
																    self.unet.config.sample_size[1] * self.mel.get_sample_rate() / self.mel.x_res / self.mel.hop_length
																)
																snake_case        =     int(mask_start_secs * pixels_per_second      )
																snake_case        =     int(mask_end_secs * pixels_per_second      )
																snake_case        =     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      ):
																						snake_case        =     self.unet(lowerCAmelCase						,      lowerCAmelCase						,      lowerCAmelCase      )['sample']
																else:
																						snake_case        =     self.unet(lowerCAmelCase						,      lowerCAmelCase      )['sample']

																if isinstance(self.scheduler						,      lowerCAmelCase      ):
																						snake_case        =     self.scheduler.step(
																						    model_output=lowerCAmelCase						,      timestep=lowerCAmelCase						,      sample=lowerCAmelCase						,      eta=lowerCAmelCase						,      generator=lowerCAmelCase						,      )['prev_sample']
																else:
																						snake_case        =     self.scheduler.step(
																						    model_output=lowerCAmelCase						,      timestep=lowerCAmelCase						,      sample=lowerCAmelCase						,      generator=lowerCAmelCase						,      )['prev_sample']

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

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

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

										snake_case        =     [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      )
										snake_case        =     np.array(
										    [np.frombuffer(image.tobytes()						,      dtype='uint8'      ).reshape((1, image.height, image.width)      ) for image in images]      )
										snake_case        =     (sample / 2_55) * 2 - 1
										snake_case        =     torch.Tensor(lowerCAmelCase      ).to(self.device      )

										for t in self.progress_bar(torch.flip(self.scheduler.timesteps						,      (0,)      )      ):
																snake_case        =     t - self.scheduler.config.num_train_timesteps // self.scheduler.num_inference_steps
																snake_case        =     self.scheduler.alphas_cumprod[t]
																snake_case        =     (
																    self.scheduler.alphas_cumprod[prev_timestep]
																    if prev_timestep >= 0
																    else self.scheduler.final_alpha_cumprod
																)
																snake_case        =     1 - alpha_prod_t
																snake_case        =     self.unet(lowerCAmelCase						,      lowerCAmelCase      )['sample']
																snake_case        =     (1 - alpha_prod_t_prev) ** 0.5 * model_output
																snake_case        =     (sample - pred_sample_direction) * alpha_prod_t_prev ** (-0.5)
																snake_case        =     sample * alpha_prod_t ** 0.5 + beta_prod_t ** 0.5 * model_output

										return sample

				@staticmethod
				def 						snake_case	(			lowerCAmelCase						,      lowerCAmelCase						,      lowerCAmelCase      ):






										"""simple docstring"""

										snake_case        =     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      )



"""simple docstring"""
import json
import os
import tempfile
import unittest
import unittest.mock as mock
from pathlib import Path

from requests.exceptions import HTTPError

from transformers.utils import (
    CONFIG_NAME,
    FLAX_WEIGHTS_NAME,
    TF2_WEIGHTS_NAME,
    TRANSFORMERS_CACHE,
    WEIGHTS_NAME,
    cached_file,
    get_file_from_repo,
    has_file,
)


SCREAMING_SNAKE_CASE__   =					"hf-internal-testing/tiny-random-bert"
SCREAMING_SNAKE_CASE__   =					os.path.join(TRANSFORMERS_CACHE, "models--hf-internal-testing--tiny-random-bert")
SCREAMING_SNAKE_CASE__   =					"9b8c223d42b2188cb49d29af482996f9d0f3e5a6"



class 					lowerCAmelCase_      (     unittest.TestCase   ):



				"""simple docstring"""

				def 						snake_case	(			self      ):






										"""simple docstring"""

										snake_case        =     cached_file(lowerCAmelCase						,      lowerCAmelCase      )
										# Should have downloaded the file in here
										self.assertTrue(os.path.isdir(lowerCAmelCase      )      )
										# Cache should contain at least those three subfolders:
										for subfolder in ["blobs", "refs", "snapshots"]:
																self.assertTrue(os.path.isdir(os.path.join(lowerCAmelCase						,      lowerCAmelCase      )      )      )
										with open(os.path.join(lowerCAmelCase						,      'refs'						,      'main'      )      ) as f:
																snake_case        =     f.read()
										self.assertEqual(lowerCAmelCase						,      os.path.join(lowerCAmelCase						,      'snapshots'						,      lowerCAmelCase						,      lowerCAmelCase      )      )
										self.assertTrue(os.path.isfile(lowerCAmelCase      )      )

										# File is cached at the same place the second time.
										snake_case        =     cached_file(lowerCAmelCase						,      lowerCAmelCase      )
										self.assertEqual(lowerCAmelCase						,      lowerCAmelCase      )

										# Using a specific revision to test the full commit hash.
										snake_case        =     cached_file(lowerCAmelCase						,      lowerCAmelCase						,      revision='9b8c223'      )
										self.assertEqual(lowerCAmelCase						,      os.path.join(lowerCAmelCase						,      'snapshots'						,      lowerCAmelCase						,      lowerCAmelCase      )      )

				def 						snake_case	(			self      ):






										"""simple docstring"""

										with self.assertRaisesRegex(lowerCAmelCase						,      'is not a valid model identifier'      ):
																snake_case        =     cached_file('tiny-random-bert'						,      lowerCAmelCase      )

										with self.assertRaisesRegex(lowerCAmelCase						,      'is not a valid git identifier'      ):
																snake_case        =     cached_file(lowerCAmelCase						,      lowerCAmelCase						,      revision='aaaa'      )

										with self.assertRaisesRegex(lowerCAmelCase						,      'does not appear to have a file named'      ):
																snake_case        =     cached_file(lowerCAmelCase						,      'conf'      )

				def 						snake_case	(			self      ):






										"""simple docstring"""

										with self.assertRaisesRegex(lowerCAmelCase						,      'does not appear to have a file named'      ):
																snake_case        =     cached_file(lowerCAmelCase						,      'conf'      )

										with open(os.path.join(lowerCAmelCase						,      'refs'						,      'main'      )      ) as f:
																snake_case        =     f.read()
										self.assertTrue(os.path.isfile(os.path.join(lowerCAmelCase						,      '.no_exist'						,      lowerCAmelCase						,      'conf'      )      )      )

										snake_case        =     cached_file(lowerCAmelCase						,      'conf'						,      _raise_exceptions_for_missing_entries=lowerCAmelCase      )
										self.assertIsNone(lowerCAmelCase      )

										snake_case        =     cached_file(lowerCAmelCase						,      'conf'						,      local_files_only=lowerCAmelCase						,      _raise_exceptions_for_missing_entries=lowerCAmelCase      )
										self.assertIsNone(lowerCAmelCase      )

										snake_case        =     mock.Mock()
										snake_case        =     5_00
										snake_case        =     {}
										snake_case        =     HTTPError
										snake_case        =     {}

										# Under the mock environment we get a 500 error when trying to reach the tokenizer.
										with mock.patch('requests.Session.request'						,      return_value=lowerCAmelCase      ) as mock_head:
																snake_case        =     cached_file(lowerCAmelCase						,      'conf'						,      _raise_exceptions_for_connection_errors=lowerCAmelCase      )
																self.assertIsNone(lowerCAmelCase      )
																# This check we did call the fake head request
																mock_head.assert_called()

				def 						snake_case	(			self      ):






										"""simple docstring"""

										self.assertTrue(has_file('hf-internal-testing/tiny-bert-pt-only'						,      lowerCAmelCase      )      )
										self.assertFalse(has_file('hf-internal-testing/tiny-bert-pt-only'						,      lowerCAmelCase      )      )
										self.assertFalse(has_file('hf-internal-testing/tiny-bert-pt-only'						,      lowerCAmelCase      )      )

				def 						snake_case	(			self      ):






										"""simple docstring"""

										self.assertIsNone(get_file_from_repo('bert-base-cased'						,      'ahah.txt'      )      )

										# The function raises if the repository does not exist.
										with self.assertRaisesRegex(lowerCAmelCase						,      'is not a valid model identifier'      ):
																get_file_from_repo('bert-base-case'						,      lowerCAmelCase      )

										# The function raises if the revision does not exist.
										with self.assertRaisesRegex(lowerCAmelCase						,      'is not a valid git identifier'      ):
																get_file_from_repo('bert-base-cased'						,      lowerCAmelCase						,      revision='ahaha'      )

										snake_case        =     get_file_from_repo('bert-base-cased'						,      lowerCAmelCase      )
										# The name is the cached name which is not very easy to test, so instead we load the content.
										snake_case        =     json.loads(open(lowerCAmelCase						,      'r'      ).read()      )
										self.assertEqual(config['hidden_size']						,      7_68      )

				def 						snake_case	(			self      ):






										"""simple docstring"""

										with tempfile.TemporaryDirectory() as tmp_dir:
																snake_case        =     Path(lowerCAmelCase      ) / 'a.txt'
																filename.touch()
																self.assertEqual(get_file_from_repo(lowerCAmelCase						,      'a.txt'      )						,      str(lowerCAmelCase      )      )

																self.assertIsNone(get_file_from_repo(lowerCAmelCase						,      'b.txt'      )      )


