

import itertools
import os
import random
import tempfile
import unittest

import numpy as np
from datasets import load_dataset

from transformers import is_speech_available
from transformers.testing_utils import check_json_file_has_correct_format, require_torch, require_torchaudio
from transformers.utils.import_utils import is_torch_available

from ...test_sequence_feature_extraction_common import SequenceFeatureExtractionTestMixin


if is_speech_available():
  from transformers import WhisperFeatureExtractor

if is_torch_available():
  import torch

__UpperCAmelCase    				=     random.Random()







def 	_snake_case		(    SCREAMING_SNAKE_CASE		,							SCREAMING_SNAKE_CASE=1.0		,							SCREAMING_SNAKE_CASE=None		,							SCREAMING_SNAKE_CASE=None				)       ->      Any:

     """simple docstring"""







     if rng is None:
          _lowerCAmelCase   :					str        =    global_rng

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

     return values





@require_torch
@require_torchaudio
class A__    ( unittest.TestCase ):



 """simple docstring"""







 def __init__(		self :	List[Any]     ,						A_ :	Optional[int]     ,						A_ :	int=7     ,						A_ :	List[Any]=4_0_0     ,						A_ :	str=2_0_0_0     ,						A_ :	Optional[int]=1_0     ,						A_ :	Any=1_6_0     ,						A_ :	int=8     ,						A_ :	Any=0.0     ,						A_ :	Optional[Any]=4_0_0_0     ,						A_ :	List[Any]=False     ,						A_ :	Optional[Any]=True     ,						):






      '''simple docstring'''




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

 def    __magic_name__       (		self :	Union[str, Any]       ):






      '''simple docstring'''




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


 def    __magic_name__       (		self :	Tuple     ,						A_ :	Optional[Any]=False     ,						A_ :	Optional[int]=False       ):






      '''simple docstring'''




      def _flatten(A_ :	Tuple       ):
           return list(itertools.chain(*A_       )       )

      if equal_length:
           _lowerCAmelCase   :					Optional[int]        =    [floats_list((self.max_seq_length, self.feature_size)       ) for _ in range(self.batch_size       )]
      else:
           # make sure that inputs increase in size
           _lowerCAmelCase   :					Dict        =    [
               floats_list((x, self.feature_size)       )
               for x in range(self.min_seq_length     ,						self.max_seq_length     ,						self.seq_length_diff       )
           ]
      if numpify:
           _lowerCAmelCase   :					int        =    [np.asarray(A_       ) for x in speech_inputs]
      return speech_inputs





@require_torch
@require_torchaudio
class A__    ( A  ,		unittest.TestCase ):



 """simple docstring"""







 _lowercase     :					Any      =						WhisperFeatureExtractor if is_speech_available() else None

 def    __magic_name__       (		self :	List[str]       ):






      '''simple docstring'''




      _lowerCAmelCase   :					Optional[int]        =    WhisperFeatureExtractionTester(self       )

 def    __magic_name__       (		self :	Dict       ):






      '''simple docstring'''




      _lowerCAmelCase   :					Any        =    self.feature_extraction_class(**self.feat_extract_dict       )

      with tempfile.TemporaryDirectory() as tmpdirname:
           _lowerCAmelCase   :					Union[str, Any]        =    feat_extract_first.save_pretrained(A_       )[0]
           check_json_file_has_correct_format(A_       )
           _lowerCAmelCase   :					str        =    self.feature_extraction_class.from_pretrained(A_       )

      _lowerCAmelCase   :					Optional[int]        =    feat_extract_first.to_dict()
      _lowerCAmelCase   :					Optional[Any]        =    feat_extract_second.to_dict()
      _lowerCAmelCase   :					int        =    feat_extract_first.mel_filters
      _lowerCAmelCase   :					Dict        =    feat_extract_second.mel_filters
      self.assertTrue(np.allclose(A_     ,						A_       )       )
      self.assertEqual(A_     ,						A_       )

 def    __magic_name__       (		self :	List[Any]       ):






      '''simple docstring'''




      _lowerCAmelCase   :					Dict        =    self.feature_extraction_class(**self.feat_extract_dict       )

      with tempfile.TemporaryDirectory() as tmpdirname:
           _lowerCAmelCase   :					Optional[Any]        =    os.path.join(A_     ,						"feat_extract.json"       )
           feat_extract_first.to_json_file(A_       )
           _lowerCAmelCase   :					Dict        =    self.feature_extraction_class.from_json_file(A_       )

      _lowerCAmelCase   :					Tuple        =    feat_extract_first.to_dict()
      _lowerCAmelCase   :					List[Any]        =    feat_extract_second.to_dict()
      _lowerCAmelCase   :					Optional[Any]        =    feat_extract_first.mel_filters
      _lowerCAmelCase   :					Tuple        =    feat_extract_second.mel_filters
      self.assertTrue(np.allclose(A_     ,						A_       )       )
      self.assertEqual(A_     ,						A_       )

 def    __magic_name__       (		self :	List[Any]       ):






      '''simple docstring'''




      _lowerCAmelCase   :					Tuple        =    self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()       )
      # create three inputs of length 800, 1000, and 1200
      _lowerCAmelCase   :					Dict        =    [floats_list((1, x)       )[0] for x in range(8_0_0     ,						1_4_0_0     ,						2_0_0       )]
      _lowerCAmelCase   :					Union[str, Any]        =    [np.asarray(A_       ) for speech_input in speech_inputs]

      # Test feature size
      _lowerCAmelCase   :					List[Any]        =    feature_extractor(A_     ,						padding="max_length"     ,						return_tensors="np"       ).input_features
      self.assertTrue(input_features.ndim == 3       )
      self.assertTrue(input_features.shape[-1] == feature_extractor.nb_max_frames       )
      self.assertTrue(input_features.shape[-2] == feature_extractor.feature_size       )

      # Test not batched input
      _lowerCAmelCase   :					Optional[Any]        =    feature_extractor(speech_inputs[0]     ,						return_tensors="np"       ).input_features
      _lowerCAmelCase   :					List[Any]        =    feature_extractor(np_speech_inputs[0]     ,						return_tensors="np"       ).input_features
      self.assertTrue(np.allclose(A_     ,						A_     ,						atol=1E-3       )       )

      # Test batched
      _lowerCAmelCase   :					Any        =    feature_extractor(A_     ,						return_tensors="np"       ).input_features
      _lowerCAmelCase   :					str        =    feature_extractor(A_     ,						return_tensors="np"       ).input_features
      for enc_seq_a, enc_seq_a in zip(A_     ,						A_       ):
           self.assertTrue(np.allclose(A_     ,						A_     ,						atol=1E-3       )       )

      # Test 2-D numpy arrays are batched.
      _lowerCAmelCase   :					str        =    [floats_list((1, x)       )[0] for x in (8_0_0, 8_0_0, 8_0_0)]
      _lowerCAmelCase   :					int        =    np.asarray(A_       )
      _lowerCAmelCase   :					str        =    feature_extractor(A_     ,						return_tensors="np"       ).input_features
      _lowerCAmelCase   :					Tuple        =    feature_extractor(A_     ,						return_tensors="np"       ).input_features
      for enc_seq_a, enc_seq_a in zip(A_     ,						A_       ):
           self.assertTrue(np.allclose(A_     ,						A_     ,						atol=1E-3       )       )

      # Test truncation required
      _lowerCAmelCase   :					Any        =    [floats_list((1, x)       )[0] for x in range(2_0_0     ,						(feature_extractor.n_samples + 5_0_0)     ,						2_0_0       )]
      _lowerCAmelCase   :					List[Any]        =    [np.asarray(A_       ) for speech_input in speech_inputs]

      _lowerCAmelCase   :					List[Any]        =    [x[: feature_extractor.n_samples] for x in speech_inputs]
      _lowerCAmelCase   :					str        =    [np.asarray(A_       ) for speech_input in speech_inputs_truncated]

      _lowerCAmelCase   :					List[str]        =    feature_extractor(A_     ,						return_tensors="np"       ).input_features
      _lowerCAmelCase   :					str        =    feature_extractor(A_     ,						return_tensors="np"       ).input_features
      for enc_seq_a, enc_seq_a in zip(A_     ,						A_       ):
           self.assertTrue(np.allclose(A_     ,						A_     ,						atol=1E-3       )       )

 def    __magic_name__       (		self :	Dict       ):






      '''simple docstring'''




      import torch

      _lowerCAmelCase   :					Optional[int]        =    self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()       )
      _lowerCAmelCase   :					int        =    np.random.rand(1_0_0     ,						3_2       ).astype(np.floataa       )
      _lowerCAmelCase   :					str        =    np_speech_inputs.tolist()

      for inputs in [py_speech_inputs, np_speech_inputs]:
           _lowerCAmelCase   :					List[Any]        =    feature_extractor.pad([{"input_features": inputs}]     ,						return_tensors="np"       )
           self.assertTrue(np_processed.input_features.dtype == np.floataa       )
           _lowerCAmelCase   :					str        =    feature_extractor.pad([{"input_features": inputs}]     ,						return_tensors="pt"       )
           self.assertTrue(pt_processed.input_features.dtype == torch.floataa       )

 def    __magic_name__       (		self :	str     ,						A_ :	int       ):






      '''simple docstring'''




      _lowerCAmelCase   :					Optional[Any]        =    load_dataset("hf-internal-testing/librispeech_asr_dummy"     ,						"clean"     ,						split="validation"       )
      # automatic decoding with librispeech
      _lowerCAmelCase   :					Any        =    ds.sort("id"       ).select(range(A_       )       )[:num_samples]["audio"]

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

 def    __magic_name__       (		self :	Dict       ):






      '''simple docstring'''




      _lowerCAmelCase   :					Any        =    torch.tensor(
          [
              0.1193, -0.0946, -0.1098, -0.0196, 0.0225, -0.0690, -0.1736, 0.0951,
              0.0971, -0.0817, -0.0702, 0.0162, 0.0260, 0.0017, -0.0192, -0.1678,
              0.0709, -0.1867, -0.0655, -0.0274, -0.0234, -0.1884, -0.0516, -0.0554,
              -0.0274, -0.1425, -0.1423, 0.0837, 0.0377, -0.0854
          ]       )
      # fmt: on

      _lowerCAmelCase   :					Tuple        =    self._load_datasamples(1       )
      _lowerCAmelCase   :					List[Any]        =    WhisperFeatureExtractor()
      _lowerCAmelCase   :					int        =    feature_extractor(A_     ,						return_tensors="pt"       ).input_features
      self.assertEqual(input_features.shape     ,						(1, 8_0, 3_0_0_0)       )
      self.assertTrue(torch.allclose(input_features[0, 0, :3_0]     ,						A_     ,						atol=1E-4       )       )


 def    __magic_name__       (		self :	Tuple       ):






      '''simple docstring'''




      _lowerCAmelCase   :					Tuple        =    self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()       )
      _lowerCAmelCase   :					Union[str, Any]        =    self._load_datasamples(1       )[0]
      _lowerCAmelCase   :					str        =    ((audio - audio.min()) / (audio.max() - audio.min())) * 6_5_5_3_5  # Rescale to [0, 65535] to show issue
      _lowerCAmelCase   :					int        =    feat_extract.zero_mean_unit_var_norm([audio]     ,						attention_mask=A_       )[0]

      self.assertTrue(np.all(np.mean(A_       ) < 1E-3       )       )
      self.assertTrue(np.all(np.abs(np.var(A_       ) - 1       ) < 1E-3       )       )









import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image

from diffusers import (
    DDIMScheduler,
    KandinskyVaaControlnetImgaImgPipeline,
    KandinskyVaaPriorEmbaEmbPipeline,
    UNetaDConditionModel,
    VQModel,
)
from diffusers.utils import floats_tensor, load_image, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()





class A__    ( A  ,		unittest.TestCase ):



 """simple docstring"""







 _lowercase     :					List[str]      =						KandinskyVaaControlnetImgaImgPipeline
 _lowercase     :					Tuple      =						['''image_embeds''', '''negative_image_embeds''', '''image''', '''hint''']
 _lowercase     :					str      =						['''image_embeds''', '''negative_image_embeds''', '''image''', '''hint''']
 _lowercase     :					Optional[Any]      =						[
     '''generator''',
     '''height''',
     '''width''',
     '''strength''',
     '''guidance_scale''',
     '''num_inference_steps''',
     '''return_dict''',
     '''guidance_scale''',
     '''num_images_per_prompt''',
     '''output_type''',
     '''return_dict''',
 ]
 _lowercase     :					Tuple      =						False

 @property
 def    __magic_name__       (		self :	Any       ):






      '''simple docstring'''




      return 3_2

 @property
 def    __magic_name__       (		self :	Tuple       ):






      '''simple docstring'''




      return 3_2

 @property
 def    __magic_name__       (		self :	List[Any]       ):






      '''simple docstring'''




      return self.time_input_dim

 @property
 def    __magic_name__       (		self :	Optional[Any]       ):






      '''simple docstring'''




      return self.time_input_dim * 4

 @property
 def    __magic_name__       (		self :	Any       ):






      '''simple docstring'''




      return 1_0_0

 @property
 def    __magic_name__       (		self :	Any       ):






      '''simple docstring'''




      torch.manual_seed(0       )

      _lowerCAmelCase   :					Tuple        =    {
          "in_channels": 8,
          # Out channels is double in channels because predicts mean and variance
          "out_channels": 8,
          "addition_embed_type": "image_hint",
          "down_block_types": ("ResnetDownsampleBlock2D", "SimpleCrossAttnDownBlock2D"),
          "up_block_types": ("SimpleCrossAttnUpBlock2D", "ResnetUpsampleBlock2D"),
          "mid_block_type": "UNetMidBlock2DSimpleCrossAttn",
          "block_out_channels": (self.block_out_channels_a, self.block_out_channels_a * 2),
          "layers_per_block": 1,
          "encoder_hid_dim": self.text_embedder_hidden_size,
          "encoder_hid_dim_type": "image_proj",
          "cross_attention_dim": self.cross_attention_dim,
          "attention_head_dim": 4,
          "resnet_time_scale_shift": "scale_shift",
          "class_embed_type": None,
      }

      _lowerCAmelCase   :					Optional[int]        =    UNetaDConditionModel(**A_       )
      return model

 @property
 def    __magic_name__       (		self :	Dict       ):






      '''simple docstring'''




      return {
          "block_out_channels": [3_2, 3_2, 6_4, 6_4],
          "down_block_types": [
              "DownEncoderBlock2D",
              "DownEncoderBlock2D",
              "DownEncoderBlock2D",
              "AttnDownEncoderBlock2D",
          ],
          "in_channels": 3,
          "latent_channels": 4,
          "layers_per_block": 1,
          "norm_num_groups": 8,
          "norm_type": "spatial",
          "num_vq_embeddings": 1_2,
          "out_channels": 3,
          "up_block_types": ["AttnUpDecoderBlock2D", "UpDecoderBlock2D", "UpDecoderBlock2D", "UpDecoderBlock2D"],
          "vq_embed_dim": 4,
      }

 @property
 def    __magic_name__       (		self :	Optional[int]       ):






      '''simple docstring'''




      torch.manual_seed(0       )
      _lowerCAmelCase   :					List[str]        =    VQModel(**self.dummy_movq_kwargs       )
      return model

 def    __magic_name__       (		self :	int       ):






      '''simple docstring'''




      _lowerCAmelCase   :					Tuple        =    self.dummy_unet
      _lowerCAmelCase   :					List[Any]        =    self.dummy_movq

      _lowerCAmelCase   :					Tuple        =    {
          "num_train_timesteps": 1_0_0_0,
          "beta_schedule": "linear",
          "beta_start": 0.00085,
          "beta_end": 0.012,
          "clip_sample": False,
          "set_alpha_to_one": False,
          "steps_offset": 0,
          "prediction_type": "epsilon",
          "thresholding": False,
      }

      _lowerCAmelCase   :					int        =    DDIMScheduler(**A_       )

      _lowerCAmelCase   :					Tuple        =    {
          "unet": unet,
          "scheduler": scheduler,
          "movq": movq,
      }

      return components

 def    __magic_name__       (		self :	Union[str, Any]     ,						A_ :	Union[str, Any]     ,						A_ :	Optional[int]=0       ):






      '''simple docstring'''




      _lowerCAmelCase   :					Tuple        =    floats_tensor((1, self.text_embedder_hidden_size)     ,						rng=random.Random(A_       )       ).to(A_       )
      _lowerCAmelCase   :					str        =    floats_tensor((1, self.text_embedder_hidden_size)     ,						rng=random.Random(seed + 1       )       ).to(
          A_       )
      # create init_image
      _lowerCAmelCase   :					Dict        =    floats_tensor((1, 3, 6_4, 6_4)     ,						rng=random.Random(A_       )       ).to(A_       )
      _lowerCAmelCase   :					Dict        =    image.cpu().permute(0     ,						2     ,						3     ,						1       )[0]
      _lowerCAmelCase   :					Tuple        =    Image.fromarray(np.uinta(A_       )       ).convert("RGB"       ).resize((2_5_6, 2_5_6)       )
      # create hint
      _lowerCAmelCase   :					Union[str, Any]        =    floats_tensor((1, 3, 6_4, 6_4)     ,						rng=random.Random(A_       )       ).to(A_       )

      if str(A_       ).startswith("mps"       ):
           _lowerCAmelCase   :					Tuple        =    torch.manual_seed(A_       )
      else:
           _lowerCAmelCase   :					Any        =    torch.Generator(device=A_       ).manual_seed(A_       )
      _lowerCAmelCase   :					Union[str, Any]        =    {
          "image": init_image,
          "image_embeds": image_embeds,
          "negative_image_embeds": negative_image_embeds,
          "hint": hint,
          "generator": generator,
          "height": 6_4,
          "width": 6_4,
          "num_inference_steps": 1_0,
          "guidance_scale": 7.0,
          "strength": 0.2,
          "output_type": "np",
      }
      return inputs


 def    __magic_name__       (		self :	Optional[int]       ):






      '''simple docstring'''




      _lowerCAmelCase   :					Union[str, Any]        =    "cpu"

      _lowerCAmelCase   :					Tuple        =    self.get_dummy_components()

      _lowerCAmelCase   :					Optional[int]        =    self.pipeline_class(**A_       )
      _lowerCAmelCase   :					Optional[Any]        =    pipe.to(A_       )

      pipe.set_progress_bar_config(disable=A_       )

      _lowerCAmelCase   :					List[str]        =    pipe(**self.get_dummy_inputs(A_       )       )
      _lowerCAmelCase   :					List[Any]        =    output.images

      _lowerCAmelCase   :					str        =    pipe(
          **self.get_dummy_inputs(A_       )     ,						return_dict=A_     ,						)[0]

      _lowerCAmelCase   :					int        =    image[0, -3:, -3:, -1]
      _lowerCAmelCase   :					List[str]        =    image_from_tuple[0, -3:, -3:, -1]

      assert image.shape == (1, 6_4, 6_4, 3)

      _lowerCAmelCase   :					List[Any]        =    np.array(
          [0.54985034, 0.55509365, 0.52561504, 0.5570494, 0.5593818, 0.5263979, 0.50285643, 0.5069846, 0.51196736]       )
      assert (
          np.abs(image_slice.flatten() - expected_slice       ).max() < 1E-2
      ), F''' expected_slice {expected_slice}, but got {image_slice.flatten()}'''
      assert (
          np.abs(image_from_tuple_slice.flatten() - expected_slice       ).max() < 1E-2
      ), F''' expected_slice {expected_slice}, but got {image_from_tuple_slice.flatten()}'''





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



 """simple docstring"""







 def    __magic_name__       (		self :	Any       ):






      '''simple docstring'''




      super().tearDown()
      gc.collect()
      torch.cuda.empty_cache()


 def    __magic_name__       (		self :	List[Any]       ):






      '''simple docstring'''




      _lowerCAmelCase   :					List[str]        =    load_numpy(
          "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
          "/kandinskyv22/kandinskyv22_controlnet_img2img_robotcat_fp16.npy"       )

      _lowerCAmelCase   :					str        =    load_image(
          "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main" "/kandinsky/cat.png"       )
      _lowerCAmelCase   :					Union[str, Any]        =    init_image.resize((5_1_2, 5_1_2)       )

      _lowerCAmelCase   :					Dict        =    load_image(
          "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
          "/kandinskyv22/hint_image_cat.png"       )
      _lowerCAmelCase   :					Tuple        =    torch.from_numpy(np.array(A_       )       ).float() / 255.0
      _lowerCAmelCase   :					Optional[int]        =    hint.permute(2     ,						0     ,						1       ).unsqueeze(0       )

      _lowerCAmelCase   :					List[str]        =    "A robot, 4k photo"

      _lowerCAmelCase   :					Dict        =    KandinskyVaaPriorEmbaEmbPipeline.from_pretrained(
          "kandinsky-community/kandinsky-2-2-prior"     ,						torch_dtype=torch.floataa       )
      pipe_prior.to(A_       )

      _lowerCAmelCase   :					str        =    KandinskyVaaControlnetImgaImgPipeline.from_pretrained(
          "kandinsky-community/kandinsky-2-2-controlnet-depth"     ,						torch_dtype=torch.floataa       )
      _lowerCAmelCase   :					Optional[int]        =    pipeline.to(A_       )

      pipeline.set_progress_bar_config(disable=A_       )

      _lowerCAmelCase   :					Dict        =    torch.Generator(device="cpu"       ).manual_seed(0       )

      _lowerCAmelCase      ,			_lowerCAmelCase   :					Tuple        =    pipe_prior(
          A_     ,						image=A_     ,						strength=0.85     ,						generator=A_     ,						negative_prompt=""     ,						).to_tuple()

      _lowerCAmelCase   :					List[Any]        =    pipeline(
          image=A_     ,						image_embeds=A_     ,						negative_image_embeds=A_     ,						hint=A_     ,						generator=A_     ,						num_inference_steps=1_0_0     ,						height=5_1_2     ,						width=5_1_2     ,						strength=0.5     ,						output_type="np"     ,						)

      _lowerCAmelCase   :					str        =    output.images[0]

      assert image.shape == (5_1_2, 5_1_2, 3)

      assert_mean_pixel_difference(A_     ,						A_       )






