






"""simple docstring"""
import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import CLIPImageProcessor, CLIPVisionConfig, CLIPVisionModelWithProjection

from diffusers import (
    AutoencoderKL,
    DPMSolverMultistepScheduler,
    PNDMScheduler,
    StableDiffusionImageVariationPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import floats_tensor, load_image, load_numpy, nightly, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import IMAGE_VARIATION_BATCH_PARAMS, IMAGE_VARIATION_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()



class 					UpperCamelCase__   (							UpperCamelCase__       ,  UpperCamelCase__       ,  UpperCamelCase__       ,  unittest.TestCase):







    """simple docstring"""
    __UpperCAmelCase											=   StableDiffusionImageVariationPipeline
    __UpperCAmelCase											=   IMAGE_VARIATION_PARAMS
    __UpperCAmelCase											=   IMAGE_VARIATION_BATCH_PARAMS
    __UpperCAmelCase											=   frozenset([])
    # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
    __UpperCAmelCase											=   frozenset([])







    def 							a__ (		self   :  List[Any]   ):





           '''simple docstring'''

           torch.manual_seed(0   )
           __magic_name__										=						UNetaDConditionModel(
               block_out_channels=(3_2, 6_4)							,						layers_per_block=2							,						sample_size=3_2							,						in_channels=4							,						out_channels=4							,						down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')							,						up_block_types=('CrossAttnUpBlock2D', 'UpBlock2D')							,						cross_attention_dim=3_2							,						)
           __magic_name__										=						PNDMScheduler(skip_prk_steps=lowerCAmelCase__   )
           torch.manual_seed(0   )
           __magic_name__										=						AutoencoderKL(
               block_out_channels=[3_2, 6_4]							,						in_channels=3							,						out_channels=3							,						down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D']							,						up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D']							,						latent_channels=4							,						)
           torch.manual_seed(0   )
           __magic_name__										=						CLIPVisionConfig(
               hidden_size=3_2							,						projection_dim=3_2							,						intermediate_size=3_7							,						layer_norm_eps=1e-0_5							,						num_attention_heads=4							,						num_hidden_layers=5							,						image_size=3_2							,						patch_size=4							,						)
           __magic_name__										=						CLIPVisionModelWithProjection(lowerCAmelCase__   )
           __magic_name__										=						CLIPImageProcessor(crop_size=3_2							,						size=3_2   )

           __magic_name__										=						{
               "unet": unet,
               "scheduler": scheduler,
               "vae": vae,
               "image_encoder": image_encoder,
               "feature_extractor": feature_extractor,
               "safety_checker": None,
           }
           return components







    def 							a__ (		self   :  List[str]							,						UpperCamelCase_   :  str							,						UpperCamelCase_   :  Optional[int]=0   ):





           '''simple docstring'''

           __magic_name__										=						floats_tensor((1, 3, 3_2, 3_2)							,						rng=random.Random(lowerCAmelCase__   )   )
           __magic_name__										=						image.cpu().permute(0							,						2							,						3							,						1   )[0]
           __magic_name__										=						Image.fromarray(np.uinta(lowerCAmelCase__   )   ).convert('RGB'   ).resize((3_2, 3_2)   )
           if str(lowerCAmelCase__   ).startswith('mps'   ):
                  __magic_name__										=						torch.manual_seed(lowerCAmelCase__   )
           else:
                  __magic_name__										=						torch.Generator(device=lowerCAmelCase__   ).manual_seed(lowerCAmelCase__   )
           __magic_name__										=						{
               "image": image,
               "generator": generator,
               "num_inference_steps": 2,
               "guidance_scale": 6.0,
               "output_type": "numpy",
           }
           return inputs







    def 							a__ (		self   :  Any   ):





           '''simple docstring'''

           __magic_name__										=						"cpu"  # ensure determinism for the device-dependent torch.Generator
           __magic_name__										=						self.get_dummy_components()
           __magic_name__										=						StableDiffusionImageVariationPipeline(**lowerCAmelCase__   )
           __magic_name__										=						sd_pipe.to(lowerCAmelCase__   )
           sd_pipe.set_progress_bar_config(disable=lowerCAmelCase__   )

           __magic_name__										=						self.get_dummy_inputs(lowerCAmelCase__   )
           __magic_name__										=						sd_pipe(**lowerCAmelCase__   ).images
           __magic_name__										=						image[0, -3:, -3:, -1]

           assert image.shape == (1, 6_4, 6_4, 3)
           __magic_name__										=						np.array([0.5_239, 0.5_723, 0.4_796, 0.5_049, 0.5_550, 0.4_685, 0.5_329, 0.4_891, 0.4_921]   )

           assert np.abs(image_slice.flatten() - expected_slice   ).max() < 1e-3







    def 							a__ (		self   :  Any   ):





           '''simple docstring'''

           __magic_name__										=						"cpu"  # ensure determinism for the device-dependent torch.Generator
           __magic_name__										=						self.get_dummy_components()
           __magic_name__										=						StableDiffusionImageVariationPipeline(**lowerCAmelCase__   )
           __magic_name__										=						sd_pipe.to(lowerCAmelCase__   )
           sd_pipe.set_progress_bar_config(disable=lowerCAmelCase__   )

           __magic_name__										=						self.get_dummy_inputs(lowerCAmelCase__   )
           __magic_name__										=						2 * [inputs["image"]]
           __magic_name__										=						sd_pipe(**lowerCAmelCase__   )

           __magic_name__										=						output.images

           __magic_name__										=						image[-1, -3:, -3:, -1]

           assert image.shape == (2, 6_4, 6_4, 3)
           __magic_name__										=						np.array([0.6_892, 0.5_637, 0.5_836, 0.5_771, 0.6_254, 0.6_409, 0.5_580, 0.5_569, 0.5_289]   )

           assert np.abs(image_slice.flatten() - expected_slice   ).max() < 1e-3







    def 							a__ (		self   :  int   ):





           '''simple docstring'''

           super().test_inference_batch_single_identical(expected_max_diff=3e-3   )



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







    """simple docstring"""







    def 							a__ (		self   :  List[Any]   ):





           '''simple docstring'''

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







    def 							a__ (		self   :  Optional[int]							,						UpperCamelCase_   :  int							,						UpperCamelCase_   :  Union[str, Any]="cpu"							,						UpperCamelCase_   :  Union[str, Any]=torch.floataa							,						UpperCamelCase_   :  Union[str, Any]=0   ):





           '''simple docstring'''

           __magic_name__										=						torch.Generator(device=lowerCAmelCase__   ).manual_seed(lowerCAmelCase__   )
           __magic_name__										=						load_image(
               'https://huggingface.co/datasets/diffusers/test-arrays/resolve/main'
               '/stable_diffusion_imgvar/input_image_vermeer.png'   )
           __magic_name__										=						np.random.RandomState(lowerCAmelCase__   ).standard_normal((1, 4, 6_4, 6_4)   )
           __magic_name__										=						torch.from_numpy(lowerCAmelCase__   ).to(device=lowerCAmelCase__							,						dtype=lowerCAmelCase__   )
           __magic_name__										=						{
               "image": init_image,
               "latents": latents,
               "generator": generator,
               "num_inference_steps": 3,
               "guidance_scale": 7.5,
               "output_type": "numpy",
           }
           return inputs







    def 							a__ (		self   :  Any   ):





           '''simple docstring'''

           __magic_name__										=						StableDiffusionImageVariationPipeline.from_pretrained(
               'lambdalabs/sd-image-variations-diffusers'							,						safety_checker=lowerCAmelCase__   )
           __magic_name__										=						sd_pipe.to(lowerCAmelCase__   )
           sd_pipe.set_progress_bar_config(disable=lowerCAmelCase__   )

           __magic_name__										=						self.get_inputs(lowerCAmelCase__   )
           __magic_name__										=						sd_pipe(**lowerCAmelCase__   ).images
           __magic_name__										=						image[0, -3:, -3:, -1].flatten()

           assert image.shape == (1, 5_1_2, 5_1_2, 3)
           __magic_name__										=						np.array([0.84_491, 0.90_789, 0.75_708, 0.78_734, 0.83_485, 0.70_099, 0.66_938, 0.68_727, 0.61_379]   )
           assert np.abs(image_slice - expected_slice   ).max() < 6e-3







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





           '''simple docstring'''

           __magic_name__										=						0

           def callback_fn(UpperCamelCase_   :  List[Any]							,						UpperCamelCase_   :  Tuple							,						UpperCamelCase_   :  Dict   ) -> None:
                  __magic_name__										=						True
                  nonlocal number_of_steps
                  number_of_steps += 1
                  if step == 1:
                         __magic_name__										=						latents.detach().cpu().numpy()
                         assert latents.shape == (1, 4, 6_4, 6_4)
                         __magic_name__										=						latents[0, -3:, -3:, -1]
                         __magic_name__										=						np.array(
                             [-0.1_621, 0.2_837, -0.7_979, -0.1_221, -1.3_057, 0.7_681, -2.1_191, 0.0_464, 1.6_309]   )

                         assert np.abs(latents_slice.flatten() - expected_slice   ).max() < 5e-2
                  elif step == 2:
                         __magic_name__										=						latents.detach().cpu().numpy()
                         assert latents.shape == (1, 4, 6_4, 6_4)
                         __magic_name__										=						latents[0, -3:, -3:, -1]
                         __magic_name__										=						np.array([0.6_299, 1.7_500, 1.1_992, -2.1_582, -1.8_994, 0.7_334, -0.7_090, 1.0_137, 1.5_273]   )

                         assert np.abs(latents_slice.flatten() - expected_slice   ).max() < 5e-2

           __magic_name__										=						False

           __magic_name__										=						StableDiffusionImageVariationPipeline.from_pretrained(
               'fusing/sd-image-variations-diffusers'							,						safety_checker=lowerCAmelCase__							,						torch_dtype=torch.floataa							,						)
           pipe.to(lowerCAmelCase__   )
           pipe.set_progress_bar_config(disable=lowerCAmelCase__   )
           pipe.enable_attention_slicing()

           __magic_name__										=						self.get_inputs(lowerCAmelCase__							,						dtype=torch.floataa   )
           pipe(**lowerCAmelCase__							,						callback=lowerCAmelCase__							,						callback_steps=1   )
           assert callback_fn.has_been_called
           assert number_of_steps == inputs["num_inference_steps"]







    def 							a__ (		self   :  str   ):





           '''simple docstring'''

           torch.cuda.empty_cache()
           torch.cuda.reset_max_memory_allocated()
           torch.cuda.reset_peak_memory_stats()

           __magic_name__										=						"fusing/sd-image-variations-diffusers"
           __magic_name__										=						StableDiffusionImageVariationPipeline.from_pretrained(
               lowerCAmelCase__							,						safety_checker=lowerCAmelCase__							,						torch_dtype=torch.floataa   )
           __magic_name__										=						pipe.to(lowerCAmelCase__   )
           pipe.set_progress_bar_config(disable=lowerCAmelCase__   )
           pipe.enable_attention_slicing(1   )
           pipe.enable_sequential_cpu_offload()

           __magic_name__										=						self.get_inputs(lowerCAmelCase__							,						dtype=torch.floataa   )
           __magic_name__										=						pipe(**lowerCAmelCase__   )

           __magic_name__										=						torch.cuda.max_memory_allocated()
           # make sure that less than 2.6 GB is allocated
           assert mem_bytes < 2.6 * 1_0**9



@nightly
@require_torch_gpu
class 					UpperCamelCase__   (							unittest.TestCase):







    """simple docstring"""







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





           '''simple docstring'''

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







    def 							a__ (		self   :  List[Any]							,						UpperCamelCase_   :  Any							,						UpperCamelCase_   :  Any="cpu"							,						UpperCamelCase_   :  Tuple=torch.floataa							,						UpperCamelCase_   :  Tuple=0   ):





           '''simple docstring'''

           __magic_name__										=						torch.Generator(device=lowerCAmelCase__   ).manual_seed(lowerCAmelCase__   )
           __magic_name__										=						load_image(
               'https://huggingface.co/datasets/diffusers/test-arrays/resolve/main'
               '/stable_diffusion_imgvar/input_image_vermeer.png'   )
           __magic_name__										=						np.random.RandomState(lowerCAmelCase__   ).standard_normal((1, 4, 6_4, 6_4)   )
           __magic_name__										=						torch.from_numpy(lowerCAmelCase__   ).to(device=lowerCAmelCase__							,						dtype=lowerCAmelCase__   )
           __magic_name__										=						{
               "image": init_image,
               "latents": latents,
               "generator": generator,
               "num_inference_steps": 5_0,
               "guidance_scale": 7.5,
               "output_type": "numpy",
           }
           return inputs







    def 							a__ (		self   :  Dict   ):





           '''simple docstring'''

           __magic_name__										=						StableDiffusionImageVariationPipeline.from_pretrained('fusing/sd-image-variations-diffusers'   )
           sd_pipe.to(lowerCAmelCase__   )
           sd_pipe.set_progress_bar_config(disable=lowerCAmelCase__   )

           __magic_name__										=						self.get_inputs(lowerCAmelCase__   )
           __magic_name__										=						sd_pipe(**lowerCAmelCase__   ).images[0]

           __magic_name__										=						load_numpy(
               'https://huggingface.co/datasets/diffusers/test-arrays/resolve/main'
               '/stable_diffusion_imgvar/lambdalabs_variations_pndm.npy'   )
           __magic_name__										=						np.abs(expected_image - image   ).max()
           assert max_diff < 1e-3







    def 							a__ (		self   :  int   ):





           '''simple docstring'''

           __magic_name__										=						StableDiffusionImageVariationPipeline.from_pretrained('fusing/sd-image-variations-diffusers'   )
           __magic_name__										=						DPMSolverMultistepScheduler.from_config(sd_pipe.scheduler.config   )
           sd_pipe.to(lowerCAmelCase__   )
           sd_pipe.set_progress_bar_config(disable=lowerCAmelCase__   )

           __magic_name__										=						self.get_inputs(lowerCAmelCase__   )
           __magic_name__										=						2_5
           __magic_name__										=						sd_pipe(**lowerCAmelCase__   ).images[0]

           __magic_name__										=						load_numpy(
               'https://huggingface.co/datasets/diffusers/test-arrays/resolve/main'
               '/stable_diffusion_imgvar/lambdalabs_variations_dpm_multi.npy'   )
           __magic_name__										=						np.abs(expected_image - image   ).max()
           assert max_diff < 1e-3









"""simple docstring"""
from typing import List, Optional, Union

import numpy as np

from ...feature_extraction_sequence_utils import SequenceFeatureExtractor
from ...feature_extraction_utils import BatchFeature
from ...utils import PaddingStrategy, TensorType, logging


snake_case   :   Tuple    		=		logging.get_logger(__name__)



class 					UpperCamelCase__   (							a_):







    """simple docstring"""
    __UpperCAmelCase											=   ["""input_values""", """attention_mask"""]







    def __init__(		self   :  Union[str, Any]							,						UpperCamelCase_   :  int=1							,						UpperCamelCase_   :  Optional[int]=1_6_0_0_0							,						UpperCamelCase_   :  List[str]=0.0							,						UpperCamelCase_   :  Optional[int]=False							,						UpperCamelCase_   :  Tuple=True							,						**UpperCamelCase_   :  Tuple							,						):





           '''simple docstring'''

           super().__init__(feature_size=UpperCamelCase_							,						sampling_rate=UpperCamelCase_							,						padding_value=UpperCamelCase_							,						**UpperCamelCase_   )
           __magic_name__										=						return_attention_mask
           __magic_name__										=						do_normalize







    @staticmethod
    def 							a__ (		UpperCamelCase_   :  List[np.ndarray]							,						UpperCamelCase_   :  List[np.ndarray]							,						UpperCamelCase_   :  float = 0.0   ):





           '''simple docstring'''

           if attention_mask is not None:
                  __magic_name__										=						np.array(UpperCamelCase_							,						np.intaa   )
                  __magic_name__										=						[]

                  for vector, length in zip(UpperCamelCase_							,						attention_mask.sum(-1   )   ):
                         __magic_name__										=						(vector - vector[:length].mean()) / np.sqrt(vector[:length].var() + 1e-7   )
                         if length < normed_slice.shape[0]:
                                __magic_name__										=						padding_value

                         normed_input_values.append(UpperCamelCase_   )
           else:
                  __magic_name__										=						[(x - x.mean()) / np.sqrt(x.var() + 1e-7   ) for x in input_values]

           return normed_input_values







    def __call__(		self   :  Any							,						UpperCamelCase_   :  Union[np.ndarray, List[float], List[np.ndarray], List[List[float]]]							,						UpperCamelCase_   :  Union[bool, str, PaddingStrategy] = False							,						UpperCamelCase_   :  Optional[int] = None							,						UpperCamelCase_   :  bool = False							,						UpperCamelCase_   :  Optional[int] = None							,						UpperCamelCase_   :  Optional[bool] = None							,						UpperCamelCase_   :  Optional[Union[str, TensorType]] = None							,						UpperCamelCase_   :  Optional[int] = None							,						**UpperCamelCase_   :  Any							,						):





           '''simple docstring'''

           if sampling_rate is not None:
                  if sampling_rate != self.sampling_rate:
                         raise ValueError(
                             f"""The model corresponding to this feature extractor: {self} was trained using a sampling rate of"""
                             f""" {self.sampling_rate}. Please make sure that the provided `raw_speech` input was sampled with"""
                             f""" {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.'   )

           __magic_name__										=						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}"""   )
           __magic_name__										=						is_batched_numpy or (
               isinstance(UpperCamelCase_							,						(list, tuple)   ) and (isinstance(raw_speech[0]							,						(np.ndarray, tuple, list)   ))
           )

           # always return batch
           if not is_batched:
                  __magic_name__										=						[raw_speech]

           # convert into correct format for padding
           __magic_name__										=						BatchFeature({'input_values': raw_speech}   )

           __magic_name__										=						self.pad(
               UpperCamelCase_							,						padding=UpperCamelCase_							,						max_length=UpperCamelCase_							,						truncation=UpperCamelCase_							,						pad_to_multiple_of=UpperCamelCase_							,						return_attention_mask=UpperCamelCase_							,						)

           # convert input values to correct format
           __magic_name__										=						padded_inputs['input_values']
           if not isinstance(input_values[0]							,						np.ndarray   ):
                  __magic_name__										=						[np.asarray(UpperCamelCase_							,						dtype=np.floataa   ) for array in input_values]
           elif (
               not isinstance(UpperCamelCase_							,						np.ndarray   )
               and isinstance(input_values[0]							,						np.ndarray   )
               and input_values[0].dtype is np.dtype(np.floataa   )
           ):
                  __magic_name__										=						[array.astype(np.floataa   ) for array in input_values]
           elif isinstance(UpperCamelCase_							,						np.ndarray   ) and input_values.dtype is np.dtype(np.floataa   ):
                  __magic_name__										=						input_values.astype(np.floataa   )

           # convert attention_mask to correct format
           __magic_name__										=						padded_inputs.get('attention_mask'   )
           if attention_mask is not None:
                  __magic_name__										=						[np.asarray(UpperCamelCase_							,						dtype=np.intaa   ) for array in attention_mask]

           # zero-mean and unit-variance normalization
           if self.do_normalize:
                  __magic_name__										=						(
                      attention_mask
                      if self._get_padding_strategies(UpperCamelCase_							,						max_length=UpperCamelCase_   ) is not PaddingStrategy.DO_NOT_PAD
                      else None
                  )
                  __magic_name__										=						self.zero_mean_unit_var_norm(
                      padded_inputs['input_values']							,						attention_mask=UpperCamelCase_							,						padding_value=self.padding_value   )

           if return_tensors is not None:
                  __magic_name__										=						padded_inputs.convert_to_tensors(UpperCamelCase_   )

           return padded_inputs