


"""simple docstring"""



import itertools
import random
import unittest

import numpy as np

from transformers import ASTFeatureExtractor
from transformers.testing_utils import require_torch, require_torchaudio
from transformers.utils.import_utils import is_torch_available

from ...test_sequence_feature_extraction_common import SequenceFeatureExtractionTestMixin


__snake_case       =       random.Random()

if is_torch_available():
  import torch







def     __lowerCAmelCase		(		lowercase    :    Tuple		,				lowercase    :    List[Any]=1.0		,				lowercase    :    List[str]=None		,				lowercase    :    Dict=None					)      ->		int:


    """simple docstring"""





    if rng is None:
        snake_case       : Union[str, Any]    =       global_rng

    snake_case       : int    =       []
    for batch_idx in range(shape[0]					):
        values.append([]					)
        for _ in range(shape[1]					):
            values[-1].append(rng.random() * scale					)

    return values







class _lowerCAmelCase							( unittest.TestCase      ):







   def __init__(      self		,      UpperCamelCase__		,      UpperCamelCase__=7		,      UpperCamelCase__=400		,      UpperCamelCase__=2000		,      UpperCamelCase__=1		,      UpperCamelCase__=0.0		,      UpperCamelCase__=1_6000		,      UpperCamelCase__=True		,      UpperCamelCase__=True		,      )      ->						Dict:

       '''simple docstring'''




       snake_case       : str    =       parent
       snake_case       : str    =       batch_size
       snake_case       : Optional[Any]    =       min_seq_length
       snake_case       : int    =       max_seq_length
       snake_case       : str    =       (self.max_seq_length - self.min_seq_length) // (self.batch_size - 1)
       snake_case       : List[str]    =       feature_size
       snake_case       : str    =       padding_value
       snake_case       : List[str]    =       sampling_rate
       snake_case       : Tuple    =       return_attention_mask
       snake_case       : str    =       do_normalize







   def 					lowerCamelCase						(      self		)      ->						Union[str, Any]:

       '''simple docstring'''




       return {
           "feature_size": self.feature_size,
           "padding_value": self.padding_value,
           "sampling_rate": self.sampling_rate,
           "return_attention_mask": self.return_attention_mask,
           "do_normalize": self.do_normalize,
       }







   def 					lowerCamelCase						(      self		,      UpperCamelCase__=False		,      UpperCamelCase__=False		)      ->						Any:

       '''simple docstring'''




       def _flatten(UpperCamelCase__		):
           return list(itertools.chain(*UpperCamelCase__		)		)

       if equal_length:
           snake_case       : str    =       floats_list((self.batch_size, self.max_seq_length)		)
       else:
           # make sure that inputs increase in size
           snake_case       : int    =       [
               _flatten(floats_list((x, self.feature_size)		)		)
               for x in range(self.min_seq_length		,      self.max_seq_length		,      self.seq_length_diff		)
           ]

       if numpify:
           snake_case       : Tuple    =       [np.asarray(UpperCamelCase__		) for x in speech_inputs]

       return speech_inputs







@require_torch
@require_torchaudio
class _lowerCAmelCase							( snake_case_       ,				unittest.TestCase      ):
   __UpperCAmelCase					:			Tuple				    =    ASTFeatureExtractor







   def 					lowerCamelCase						(      self		)      ->						Dict:

       '''simple docstring'''




       snake_case       : str    =       ASTFeatureExtractionTester(self		)







   def 					lowerCamelCase						(      self		)      ->						Optional[int]:

       '''simple docstring'''




       snake_case       : List[str]    =       self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()		)
       # create three inputs of length 800, 1000, and 1200
       snake_case       : str    =       [floats_list((1, x)		)[0] for x in range(800		,      1400		,      200		)]
       snake_case       : Any    =       [np.asarray(UpperCamelCase__		) for speech_input in speech_inputs]

       # Test not batched input
       snake_case       : str    =       feat_extract(speech_inputs[0]		,      return_tensors="np"		).input_values
       snake_case       : Union[str, Any]    =       feat_extract(np_speech_inputs[0]		,      return_tensors="np"		).input_values
       self.assertTrue(np.allclose(UpperCamelCase__		,      UpperCamelCase__		,      atol=1e-3		)		)

       # Test batched
       snake_case       : List[str]    =       feat_extract(UpperCamelCase__		,      padding=UpperCamelCase__		,      return_tensors="np"		).input_values
       snake_case       : List[str]    =       feat_extract(UpperCamelCase__		,      padding=UpperCamelCase__		,      return_tensors="np"		).input_values
       for enc_seq_a, enc_seq_a in zip(UpperCamelCase__		,      UpperCamelCase__		):
           self.assertTrue(np.allclose(UpperCamelCase__		,      UpperCamelCase__		,      atol=1e-3		)		)

       # Test 2-D numpy arrays are batched.
       snake_case       : Optional[int]    =       [floats_list((1, x)		)[0] for x in (800, 800, 800)]
       snake_case       : int    =       np.asarray(UpperCamelCase__		)
       snake_case       : Optional[int]    =       feat_extract(UpperCamelCase__		,      return_tensors="np"		).input_values
       snake_case       : Dict    =       feat_extract(UpperCamelCase__		,      return_tensors="np"		).input_values
       for enc_seq_a, enc_seq_a in zip(UpperCamelCase__		,      UpperCamelCase__		):
           self.assertTrue(np.allclose(UpperCamelCase__		,      UpperCamelCase__		,      atol=1e-3		)		)







   @require_torch
   def 					lowerCamelCase						(      self		)      ->						Dict:

       '''simple docstring'''




       import torch

       snake_case       : Tuple    =       self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()		)
       snake_case       : str    =       np.random.rand(100		).astype(np.floataa		)
       snake_case       : int    =       np_speech_inputs.tolist()

       for inputs in [py_speech_inputs, np_speech_inputs]:
           snake_case       : Optional[int]    =       feature_extractor.pad([{"input_values": inputs}]		,      return_tensors="np"		)
           self.assertTrue(np_processed.input_values.dtype == np.floataa		)
           snake_case       : Optional[int]    =       feature_extractor.pad([{"input_values": inputs}]		,      return_tensors="pt"		)
           self.assertTrue(pt_processed.input_values.dtype == torch.floataa		)







   def 					lowerCamelCase						(      self		,      UpperCamelCase__		)      ->						List[str]:

       '''simple docstring'''




       from datasets import load_dataset

       snake_case       : List[str]    =       load_dataset("hf-internal-testing/librispeech_asr_dummy"		,      "clean"		,      split="validation"		)
       # automatic decoding with librispeech
       snake_case       : List[str]    =       ds.sort("id"		).select(range(UpperCamelCase__		)		)[:num_samples]["audio"]

       return [x["array"] for x in speech_samples]







   @require_torch
   def 					lowerCamelCase						(      self		)      ->						Union[str, Any]:

       '''simple docstring'''




       snake_case       : Any    =       torch.tensor(
           [-0.9894, -1.2776, -0.9066, -1.2776, -0.9349, -1.2609, -1.0386, -1.2776,
            -1.1561, -1.2776, -1.2052, -1.2723, -1.2190, -1.2132, -1.2776, -1.1133,
            -1.1953, -1.1343, -1.1584, -1.2203, -1.1770, -1.2474, -1.2381, -1.1936,
            -0.9270, -0.8317, -0.8049, -0.7706, -0.7565, -0.7869]		)
       # fmt: on

       snake_case       : Optional[Any]    =       self._load_datasamples(1		)
       snake_case       : List[str]    =       ASTFeatureExtractor()
       snake_case       : Tuple    =       feature_extractor(UpperCamelCase__		,      return_tensors="pt"		).input_values
       self.assertEquals(input_values.shape		,      (1, 1024, 128)		)
       self.assertTrue(torch.allclose(input_values[0, 0, :30]		,      UpperCamelCase__		,      atol=1e-4		)		)









"""simple docstring"""



import inspect
from typing import List, Optional, Tuple, Union

import numpy as np
import PIL
import torch
import torch.utils.checkpoint

from ...models import UNetaDModel, VQModel
from ...schedulers import (
    DDIMScheduler,
    DPMSolverMultistepScheduler,
    EulerAncestralDiscreteScheduler,
    EulerDiscreteScheduler,
    LMSDiscreteScheduler,
    PNDMScheduler,
)
from ...utils import PIL_INTERPOLATION, randn_tensor
from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput







def     __lowerCAmelCase		(		lowercase    :    Optional[int]					)      ->		Union[str, Any]:


    """simple docstring"""





    snake_case   ,snake_case       : List[str]    =       image.size
    snake_case   ,snake_case       : int    =       (x - x % 32 for x in (w, h))  # resize to integer multiple of 32
    snake_case       : Any    =       image.resize((w, h)		,				resample=PIL_INTERPOLATION["lanczos"]					)
    snake_case       : Dict    =       np.array(lowercase					).astype(np.floataa					) / 255.0
    snake_case       : List[str]    =       image[None].transpose(0		,				3		,				1		,				2					)
    snake_case       : List[Any]    =       torch.from_numpy(lowercase					)
    return 2.0 * image - 1.0







class _lowerCAmelCase							( snake_case_      ):







   def __init__(      self		,      UpperCamelCase__		,      UpperCamelCase__		,      UpperCamelCase__		,      )      ->						Optional[Any]:

       '''simple docstring'''




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







   @torch.no_grad()
   def __call__(      self		,      UpperCamelCase__ = None		,      UpperCamelCase__ = 1		,      UpperCamelCase__ = 100		,      UpperCamelCase__ = 0.0		,      UpperCamelCase__ = None		,      UpperCamelCase__ = "pil"		,      UpperCamelCase__ = True		,      )      ->						Union[Tuple, ImagePipelineOutput]:

       '''simple docstring'''




       if isinstance(UpperCamelCase__		,      PIL.Image.Image		):
           snake_case       : Tuple    =       1
       elif isinstance(UpperCamelCase__		,      torch.Tensor		):
           snake_case       : str    =       image.shape[0]
       else:
           raise ValueError(F'`image` has to be of type `PIL.Image.Image` or `torch.Tensor` but is {type(UpperCamelCase__		)}'		)

       if isinstance(UpperCamelCase__		,      PIL.Image.Image		):
           snake_case       : Optional[int]    =       preprocess(UpperCamelCase__		)

       snake_case   ,snake_case       : List[str]    =       image.shape[-2:]

       # in_channels should be 6: 3 for latents, 3 for low resolution image
       snake_case       : Any    =       (batch_size, self.unet.config.in_channels // 2, height, width)
       snake_case       : List[str]    =       next(self.unet.parameters()		).dtype

       snake_case       : int    =       randn_tensor(UpperCamelCase__		,      generator=UpperCamelCase__		,      device=self.device		,      dtype=UpperCamelCase__		)

       snake_case       : List[Any]    =       image.to(device=self.device		,      dtype=UpperCamelCase__		)

       # set timesteps and move to the correct device
       self.scheduler.set_timesteps(UpperCamelCase__		,      device=self.device		)
       snake_case       : Optional[Any]    =       self.scheduler.timesteps

       # scale the initial noise by the standard deviation required by the scheduler
       snake_case       : Optional[int]    =       latents * self.scheduler.init_noise_sigma

       # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature.
       # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
       # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
       # and should be between [0, 1]
       snake_case       : Optional[int]    =       "eta" in set(inspect.signature(self.scheduler.step		).parameters.keys()		)
       snake_case       : Union[str, Any]    =       {}
       if accepts_eta:
           snake_case       : List[str]    =       eta

       for t in self.progress_bar(UpperCamelCase__		):
           # concat latents and low resolution image in the channel dimension.
           snake_case       : int    =       torch.cat([latents, image]		,      dim=1		)
           snake_case       : List[Any]    =       self.scheduler.scale_model_input(UpperCamelCase__		,      UpperCamelCase__		)
           # predict the noise residual
           snake_case       : int    =       self.unet(UpperCamelCase__		,      UpperCamelCase__		).sample
           # compute the previous noisy sample x_t -> x_t-1
           snake_case       : List[Any]    =       self.scheduler.step(UpperCamelCase__		,      UpperCamelCase__		,      UpperCamelCase__		,      **UpperCamelCase__		).prev_sample

       # decode the image latents with the VQVAE
       snake_case       : Optional[int]    =       self.vqvae.decode(UpperCamelCase__		).sample
       snake_case       : Tuple    =       torch.clamp(UpperCamelCase__		,      -1.0		,      1.0		)
       snake_case       : Union[str, Any]    =       image / 2 + 0.5
       snake_case       : Optional[Any]    =       image.cpu().permute(0		,      2		,      3		,      1		).numpy()

       if output_type == "pil":
           snake_case       : Union[str, Any]    =       self.numpy_to_pil(UpperCamelCase__		)

       if not return_dict:
           return (image,)

       return ImagePipelineOutput(images=UpperCamelCase__		)





