"""simple docstring"""


import gc
import random
import unittest

import numpy as np
import torch
from transformers import (
    CLIPImageProcessor,
    CLIPTextConfig,
    CLIPTextModelWithProjection,
    CLIPTokenizer,
    CLIPVisionConfig,
    CLIPVisionModelWithProjection,
)

from diffusers import (
    DiffusionPipeline,
    UnCLIPImageVariationPipeline,
    UnCLIPScheduler,
    UNetaDConditionModel,
    UNetaDModel,
)
from diffusers.pipelines.unclip.text_proj import UnCLIPTextProjModel
from diffusers.utils import floats_tensor, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, load_image, require_torch_gpu, skip_mps

from ..pipeline_params import IMAGE_VARIATION_BATCH_PARAMS, IMAGE_VARIATION_PARAMS
from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()






class    a__    (   _a					,					unittest.TestCase	):
   snake_case_							=      UnCLIPImageVariationPipeline
   snake_case_							=      IMAGE_VARIATION_PARAMS - {"height", "width", "guidance_scale"}
   snake_case_							=      IMAGE_VARIATION_BATCH_PARAMS

   snake_case_							=      [
       "generator",
       "return_dict",
       "decoder_num_inference_steps",
       "super_res_num_inference_steps",
   ]
   snake_case_							=      False


   @property
   def 	snake_case__	(     self  ):
        '''simple docstring'''



        return 32


   @property
   def 	snake_case__	(     self  ):
        '''simple docstring'''



        return 32


   @property
   def 	snake_case__	(     self  ):
        '''simple docstring'''



        return self.time_input_dim


   @property
   def 	snake_case__	(     self  ):
        '''simple docstring'''



        return self.time_input_dim * 4


   @property
   def 	snake_case__	(     self  ):
        '''simple docstring'''



        return 100


   @property
   def 	snake_case__	(     self  ):
        '''simple docstring'''



        lowercase__				 =							CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"  )
        return tokenizer


   @property
   def 	snake_case__	(     self  ):
        '''simple docstring'''



        torch.manual_seed(0  )
        lowercase__				 =							CLIPTextConfig(
            bos_token_id=0,				eos_token_id=2,				hidden_size=self.text_embedder_hidden_size,				projection_dim=self.text_embedder_hidden_size,				intermediate_size=37,				layer_norm_eps=1E-05,				num_attention_heads=4,				num_hidden_layers=5,				pad_token_id=1,				vocab_size=1000,				)
        return CLIPTextModelWithProjection(__a  )


   @property
   def 	snake_case__	(     self  ):
        '''simple docstring'''



        torch.manual_seed(0  )
        lowercase__				 =							CLIPVisionConfig(
            hidden_size=self.text_embedder_hidden_size,				projection_dim=self.text_embedder_hidden_size,				num_hidden_layers=5,				num_attention_heads=4,				image_size=32,				intermediate_size=37,				patch_size=1,				)
        return CLIPVisionModelWithProjection(__a  )


   @property
   def 	snake_case__	(     self  ):
        '''simple docstring'''



        torch.manual_seed(0  )

        lowercase__				 =							{
            "clip_embeddings_dim": self.text_embedder_hidden_size,
            "time_embed_dim": self.time_embed_dim,
            "cross_attention_dim": self.cross_attention_dim,
        }

        lowercase__				 =							UnCLIPTextProjModel(**__a  )
        return model


   @property
   def 	snake_case__	(     self  ):
        '''simple docstring'''



        torch.manual_seed(0  )

        lowercase__				 =							{
            "sample_size": 32,
            # RGB in channels
            "in_channels": 3,
            # Out channels is double in channels because predicts mean and variance
            "out_channels": 6,
            "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,
            "cross_attention_dim": self.cross_attention_dim,
            "attention_head_dim": 4,
            "resnet_time_scale_shift": "scale_shift",
            "class_embed_type": "identity",
        }

        lowercase__				 =							UNetaDConditionModel(**__a  )
        return model


   @property
   def 	snake_case__	(     self  ):
        '''simple docstring'''



        return {
            "sample_size": 64,
            "layers_per_block": 1,
            "down_block_types": ("ResnetDownsampleBlock2D", "ResnetDownsampleBlock2D"),
            "up_block_types": ("ResnetUpsampleBlock2D", "ResnetUpsampleBlock2D"),
            "block_out_channels": (self.block_out_channels_a, self.block_out_channels_a * 2),
            "in_channels": 6,
            "out_channels": 3,
        }


   @property
   def 	snake_case__	(     self  ):
        '''simple docstring'''



        torch.manual_seed(0  )

        lowercase__				 =							UNetaDModel(**self.dummy_super_res_kwargs  )
        return model


   @property
   def 	snake_case__	(     self  ):
        '''simple docstring'''



        torch.manual_seed(1  )

        lowercase__				 =							UNetaDModel(**self.dummy_super_res_kwargs  )
        return model


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



        lowercase__				 =							self.dummy_decoder
        lowercase__				 =							self.dummy_text_proj
        lowercase__				 =							self.dummy_text_encoder
        lowercase__				 =							self.dummy_tokenizer
        lowercase__				 =							self.dummy_super_res_first
        lowercase__				 =							self.dummy_super_res_last

        lowercase__				 =							UnCLIPScheduler(
            variance_type="learned_range",				prediction_type="epsilon",				num_train_timesteps=1000,				)

        lowercase__				 =							UnCLIPScheduler(
            variance_type="fixed_small_log",				prediction_type="epsilon",				num_train_timesteps=1000,				)

        lowercase__				 =							CLIPImageProcessor(crop_size=32,				size=32  )

        lowercase__				 =							self.dummy_image_encoder

        return {
            "decoder": decoder,
            "text_encoder": text_encoder,
            "tokenizer": tokenizer,
            "text_proj": text_proj,
            "feature_extractor": feature_extractor,
            "image_encoder": image_encoder,
            "super_res_first": super_res_first,
            "super_res_last": super_res_last,
            "decoder_scheduler": decoder_scheduler,
            "super_res_scheduler": super_res_scheduler,
        }


   def 	snake_case__	(     self,				_UpperCAmelCase,				_UpperCAmelCase=0,				_UpperCAmelCase=True  ):
        '''simple docstring'''



        lowercase__				 =							floats_tensor((1, 3, 32, 32),				rng=random.Random(__a  )  ).to(__a  )
        if str(__a  ).startswith("mps"  ):
             lowercase__				 =							torch.manual_seed(__a  )
        else:
             lowercase__				 =							torch.Generator(device=__a  ).manual_seed(__a  )

        if pil_image:
             lowercase__				 =							input_image * 0.5 + 0.5
             lowercase__				 =							input_image.clamp(0,				1  )
             lowercase__				 =							input_image.cpu().permute(0,				2,				3,				1  ).float().numpy()
             lowercase__				 =							DiffusionPipeline.numpy_to_pil(__a  )[0]

        return {
            "image": input_image,
            "generator": generator,
            "decoder_num_inference_steps": 2,
            "super_res_num_inference_steps": 2,
            "output_type": "np",
        }


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



        lowercase__				 =							"cpu"

        lowercase__				 =							self.get_dummy_components()

        lowercase__				 =							self.pipeline_class(**__a  )
        lowercase__				 =							pipe.to(__a  )

        pipe.set_progress_bar_config(disable=__a  )

        lowercase__				 =							self.get_dummy_inputs(__a,				pil_image=__a  )

        lowercase__				 =							pipe(**__a  )
        lowercase__				 =							output.images

        lowercase__				 =							self.get_dummy_inputs(__a,				pil_image=__a  )

        lowercase__				 =							pipe(
            **__a,				return_dict=__a,				)[0]

        lowercase__				 =							image[0, -3:, -3:, -1]
        lowercase__				 =							image_from_tuple[0, -3:, -3:, -1]

        assert image.shape == (1, 64, 64, 3)

        lowercase__				 =							np.array(
            [
                0.9_997,
                0.0_002,
                0.9_997,
                0.9_997,
                0.9_969,
                0.0_023,
                0.9_997,
                0.9_969,
                0.9_970,
            ]  )

        assert np.abs(image_slice.flatten() - expected_slice  ).max() < 1E-2
        assert np.abs(image_from_tuple_slice.flatten() - expected_slice  ).max() < 1E-2


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



        lowercase__				 =							"cpu"

        lowercase__				 =							self.get_dummy_components()

        lowercase__				 =							self.pipeline_class(**__a  )
        lowercase__				 =							pipe.to(__a  )

        pipe.set_progress_bar_config(disable=__a  )

        lowercase__				 =							self.get_dummy_inputs(__a,				pil_image=__a  )

        lowercase__				 =							pipe(**__a  )
        lowercase__				 =							output.images

        lowercase__				 =							self.get_dummy_inputs(__a,				pil_image=__a  )

        lowercase__				 =							pipe(
            **__a,				return_dict=__a,				)[0]

        lowercase__				 =							image[0, -3:, -3:, -1]
        lowercase__				 =							image_from_tuple[0, -3:, -3:, -1]

        assert image.shape == (1, 64, 64, 3)

        lowercase__				 =							np.array([0.9_997, 0.0_003, 0.9_997, 0.9_997, 0.9_970, 0.0_024, 0.9_997, 0.9_971, 0.9_971]  )

        assert np.abs(image_slice.flatten() - expected_slice  ).max() < 1E-2
        assert np.abs(image_from_tuple_slice.flatten() - expected_slice  ).max() < 1E-2


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



        lowercase__				 =							"cpu"

        lowercase__				 =							self.get_dummy_components()

        lowercase__				 =							self.pipeline_class(**__a  )
        lowercase__				 =							pipe.to(__a  )

        pipe.set_progress_bar_config(disable=__a  )

        lowercase__				 =							self.get_dummy_inputs(__a,				pil_image=__a  )
        lowercase__				 =							[
            pipeline_inputs["image"],
            pipeline_inputs["image"],
        ]

        lowercase__				 =							pipe(**__a  )
        lowercase__				 =							output.images

        lowercase__				 =							self.get_dummy_inputs(__a,				pil_image=__a  )
        lowercase__				 =							[
            tuple_pipeline_inputs["image"],
            tuple_pipeline_inputs["image"],
        ]

        lowercase__				 =							pipe(
            **__a,				return_dict=__a,				)[0]

        lowercase__				 =							image[0, -3:, -3:, -1]
        lowercase__				 =							image_from_tuple[0, -3:, -3:, -1]

        assert image.shape == (2, 64, 64, 3)

        lowercase__				 =							np.array(
            [
                0.9_997,
                0.9_989,
                0.0_008,
                0.0_021,
                0.9_960,
                0.0_018,
                0.0_014,
                0.0_002,
                0.9_933,
            ]  )

        assert np.abs(image_slice.flatten() - expected_slice  ).max() < 1E-2
        assert np.abs(image_from_tuple_slice.flatten() - expected_slice  ).max() < 1E-2


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



        lowercase__				 =							torch.device("cpu"  )






        class    a__    :
           snake_case_							=      1



        lowercase__				 =							self.get_dummy_components()

        lowercase__				 =							self.pipeline_class(**__a  )
        lowercase__				 =							pipe.to(__a  )

        pipe.set_progress_bar_config(disable=__a  )

        lowercase__				 =							torch.Generator(device=__a  ).manual_seed(0  )
        lowercase__				 =							pipe.decoder.dtype
        lowercase__				 =							1

        lowercase__				 =							(
            batch_size,
            pipe.decoder.config.in_channels,
            pipe.decoder.config.sample_size,
            pipe.decoder.config.sample_size,
        )
        lowercase__				 =							pipe.prepare_latents(
            __a,				dtype=__a,				device=__a,				generator=__a,				latents=__a,				scheduler=DummyScheduler()  )

        lowercase__				 =							(
            batch_size,
            pipe.super_res_first.config.in_channels // 2,
            pipe.super_res_first.config.sample_size,
            pipe.super_res_first.config.sample_size,
        )
        lowercase__				 =							pipe.prepare_latents(
            __a,				dtype=__a,				device=__a,				generator=__a,				latents=__a,				scheduler=DummyScheduler()  )

        lowercase__				 =							self.get_dummy_inputs(__a,				pil_image=__a  )

        lowercase__				 =							pipe(
            **__a,				decoder_latents=__a,				super_res_latents=__a  ).images

        lowercase__				 =							self.get_dummy_inputs(__a,				pil_image=__a  )
        # Don't pass image, instead pass embedding
        lowercase__				 =							pipeline_inputs.pop("image"  )
        lowercase__				 =							pipe.image_encoder(__a  ).image_embeds

        lowercase__				 =							pipe(
            **__a,				decoder_latents=__a,				super_res_latents=__a,				image_embeddings=__a,				).images

        # make sure passing text embeddings manually is identical
        assert np.abs(img_out_a - img_out_a  ).max() < 1E-4


   @skip_mps
   def 	snake_case__	(     self  ):
        '''simple docstring'''



        lowercase__				 =							torch_device == "cpu"

        # Check is relaxed because there is not a torch 2.0 sliced attention added kv processor
        lowercase__				 =							1E-2

        self._test_attention_slicing_forward_pass(
            test_max_difference=__a,				expected_max_diff=__a  )


   @skip_mps
   def 	snake_case__	(     self  ):
        '''simple docstring'''



        lowercase__				 =							torch_device == "cpu"
        lowercase__				 =							True
        lowercase__				 =							[
            "decoder_num_inference_steps",
            "super_res_num_inference_steps",
        ]

        self._test_inference_batch_single_identical(
            test_max_difference=__a,				relax_max_difference=__a,				additional_params_copy_to_batched_inputs=__a,				)


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



        lowercase__				 =							[
            "decoder_num_inference_steps",
            "super_res_num_inference_steps",
        ]

        if torch_device == "mps":
             # TODO: MPS errors with larger batch sizes
             lowercase__				 =							[2, 3]
             self._test_inference_batch_consistent(
                 batch_sizes=__a,				additional_params_copy_to_batched_inputs=__a,				)
        else:
             self._test_inference_batch_consistent(
                 additional_params_copy_to_batched_inputs=__a  )


   @skip_mps
   def 	snake_case__	(     self  ):
        '''simple docstring'''



        return super().test_dict_tuple_outputs_equivalent()


   @skip_mps
   def 	snake_case__	(     self  ):
        '''simple docstring'''



        return super().test_save_load_local()


   @skip_mps
   def 	snake_case__	(     self  ):
        '''simple docstring'''



        return super().test_save_load_optional_components()






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


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



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


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



        lowercase__				 =							load_image(
            "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/unclip/cat.png"  )
        lowercase__				 =							load_numpy(
            "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
            "/unclip/karlo_v1_alpha_cat_variation_fp16.npy"  )

        lowercase__				 =							UnCLIPImageVariationPipeline.from_pretrained(
            "kakaobrain/karlo-v1-alpha-image-variations",				torch_dtype=torch.floataa  )
        lowercase__				 =							pipeline.to(__a  )
        pipeline.set_progress_bar_config(disable=__a  )

        lowercase__				 =							torch.Generator(device="cpu"  ).manual_seed(0  )
        lowercase__				 =							pipeline(
            __a,				generator=__a,				output_type="np",				)

        lowercase__				 =							output.images[0]

        assert image.shape == (256, 256, 3)

        assert_mean_pixel_difference(__a,				__a,				15  )






"""simple docstring"""


from math import ceil
from typing import List, Optional, Union

import numpy as np

from ...audio_utils import mel_filter_bank, spectrogram, window_function
from ...feature_extraction_sequence_utils import BatchFeature, SequenceFeatureExtractor
from ...utils import TensorType, logging


lowerCAmelCase_:					Union[str, Any]     =					logging.get_logger(__name__)






class    a__    (   _a	):
   snake_case_							=      ["audio_values", "audio_mask"]


   def __init__(     self,				_UpperCAmelCase=2048,				_UpperCAmelCase=1,				_UpperCAmelCase=[16, 16],				_UpperCAmelCase=128,				_UpperCAmelCase=4_4100,				_UpperCAmelCase=86,				_UpperCAmelCase=2048,				_UpperCAmelCase=0.0,				**_UpperCAmelCase,				):
        '''simple docstring'''



        super().__init__(
            feature_size=_UpperCAmelCase,				sampling_rate=_UpperCAmelCase,				padding_value=_UpperCAmelCase,				**_UpperCAmelCase,				)

        lowercase__				 =							spectrogram_length
        lowercase__				 =							num_channels
        lowercase__				 =							patch_size
        lowercase__				 =							feature_size // self.patch_size[1]
        lowercase__				 =							n_fft
        lowercase__				 =							sampling_rate // hop_length_to_sampling_rate
        lowercase__				 =							sampling_rate
        lowercase__				 =							padding_value
        lowercase__				 =							mel_filter_bank(
            num_frequency_bins=1 + n_fft // 2,				num_mel_filters=_UpperCAmelCase,				min_frequency=0.0,				max_frequency=22_050.0,				sampling_rate=_UpperCAmelCase,				norm="slaney",				mel_scale="slaney",				).T


   def 	snake_case__	(     self,				_UpperCAmelCase  ):
        '''simple docstring'''



        lowercase__				 =							spectrogram(
            _UpperCAmelCase,				window_function(self.n_fft,				"hann"  ),				frame_length=self.n_fft,				hop_length=self.hop_length,				power=2.0,				mel_filters=self.mel_filters.T,				log_mel="dB",				db_range=80.0,				)
        lowercase__				 =							log_spec[:, :-1]
        lowercase__				 =							log_spec - 20.0
        lowercase__				 =							np.clip(log_spec / 40.0,				-2.0,				0.0  ) + 1.0
        return log_spec


   def __call__(     self,				_UpperCAmelCase,				_UpperCAmelCase = None,				_UpperCAmelCase = True,				_UpperCAmelCase = None,				_UpperCAmelCase = False,				_UpperCAmelCase = False,				**_UpperCAmelCase,				):
        '''simple docstring'''



        if sampling_rate is not None:
             if sampling_rate != self.sampling_rate:
                  raise ValueError(
                      "This feature extractor is set to support sampling rate"
                      F''' of {self.sampling_rate}. Please make sure that the provided `raw_speech` input was sampled'''
                      F''' with {self.sampling_rate} and not {sampling_rate}.'''  )
        else:
             logger.warning(
                 "It is strongly recommended to pass the `sampling_rate` argument to this function. "
                 "Failing to do so can result in silent errors that might be hard to debug."  )

        lowercase__				 =							isinstance(_UpperCAmelCase,				np.ndarray  ) and len(raw_speech.shape  ) > 1
        if is_batched_numpy and len(raw_speech.shape  ) > 2:
             raise ValueError(F'''Only mono-channel audio is supported for input to {self}'''  )
        lowercase__				 =							is_batched_numpy or (
            isinstance(_UpperCAmelCase,				(list, tuple)  ) and (isinstance(raw_speech[0],				(np.ndarray, tuple, list)  ))
        )
        if is_batched:
             lowercase__				 =							[np.asarray([speech],				dtype=np.floataa  ).T for speech in raw_speech]
        elif not is_batched and not isinstance(_UpperCAmelCase,				np.ndarray  ):
             lowercase__				 =							np.asarray(_UpperCAmelCase,				dtype=np.floataa  )
        elif isinstance(_UpperCAmelCase,				np.ndarray  ) and raw_speech.dtype is np.dtype(np.floataa  ):
             lowercase__				 =							raw_speech.astype(np.floataa  )
        # always return batch
        if not is_batched:
             lowercase__				 =							[np.asarray([raw_speech]  ).T]

        # Convert audio signals to log mel spectrograms, truncate by time axis
        lowercase__				 =							[
            self._np_extract_fbank_features(waveform.squeeze()  ).T[: self.spectrogram_length] for waveform in raw_speech
        ]
        if isinstance(audio_features[0],				_UpperCAmelCase  ):
             lowercase__				 =							[np.asarray(_UpperCAmelCase,				dtype=np.floataa  ) for feature in audio_features]

        # Create audio attention mask
        lowercase__				 =							max(
            [ceil(feature.shape[0] / self.patch_size[0]  ) * self.freq_len for feature in audio_features]  )  # The maximum number of audio patches in a batch
        if return_attention_mask:
             lowercase__				 =							[
                 (ceil(feature.shape[0] / self.patch_size[0]  ) * self.freq_len) * [1]
                 + (max_patch_len - ceil(feature.shape[0] / self.patch_size[0]  ) * self.freq_len) * [0]
                 for feature in audio_features
             ]
             lowercase__				 =							np.array(_UpperCAmelCase  ).astype(np.floataa  )

        # convert into correct format for padding
        lowercase__				 =							max_patch_len // self.freq_len * self.patch_size[0]  # The maximum audio size in a batch
        lowercase__				 =							np.ones([len(_UpperCAmelCase  ), 1, max_time_len, self.feature_size]  ).astype(np.floataa  )
        lowercase__				 =							padded_audio_features * self.padding_value
        for i in range(len(_UpperCAmelCase  )  ):
             lowercase__				 =							audio_features[i]
             lowercase__				 =							feature

        # return as BatchFeature
        if return_attention_mask:
             lowercase__				 =							{"audio_values": padded_audio_features, "audio_mask": audio_mask}
        else:
             lowercase__				 =							{"audio_values": padded_audio_features}

        lowercase__				 =							BatchFeature(data=_UpperCAmelCase,				tensor_type=_UpperCAmelCase  )
        return encoded_inputs




