
'''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 _A  (    __SCREAMING_SNAKE_CASE	,				__SCREAMING_SNAKE_CASE	,				__SCREAMING_SNAKE_CASE	,				unittest.TestCase    ):
		_SCREAMING_SNAKE_CASE					:    List[Any]						      =			StableDiffusionImageVariationPipeline
		_SCREAMING_SNAKE_CASE					:    int						      =			IMAGE_VARIATION_PARAMS
		_SCREAMING_SNAKE_CASE					:    Optional[int]						      =			IMAGE_VARIATION_BATCH_PARAMS
		_SCREAMING_SNAKE_CASE					:    Tuple						      =			frozenset([]    )
		# TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
		_SCREAMING_SNAKE_CASE					:    int						      =			frozenset([]    )





		def 							__A       (     self				)       ->     List[str]:

					'''simple docstring'''
					torch.manual_seed(0				)
					__UpperCAmelCase   :      Any										=     UNetaDConditionModel(
					    block_out_channels=(32, 64) ,       layers_per_block=2 ,       sample_size=32 ,       in_channels=4 ,       out_channels=4 ,       down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""") ,       up_block_types=("""CrossAttnUpBlock2D""", """UpBlock2D""") ,       cross_attention_dim=32 ,       )
					__UpperCAmelCase   :      Optional[Any]										=     PNDMScheduler(skip_prk_steps=UpperCAmelCase_				)
					torch.manual_seed(0				)
					__UpperCAmelCase   :      List[Any]										=     AutoencoderKL(
					    block_out_channels=[32, 64] ,       in_channels=3 ,       out_channels=3 ,       down_block_types=["""DownEncoderBlock2D""", """DownEncoderBlock2D"""] ,       up_block_types=["""UpDecoderBlock2D""", """UpDecoderBlock2D"""] ,       latent_channels=4 ,       )
					torch.manual_seed(0				)
					__UpperCAmelCase   :      List[Any]										=     CLIPVisionConfig(
					    hidden_size=32 ,       projection_dim=32 ,       intermediate_size=37 ,       layer_norm_eps=1E-05 ,       num_attention_heads=4 ,       num_hidden_layers=5 ,       image_size=32 ,       patch_size=4 ,       )
					__UpperCAmelCase   :      Optional[Any]										=     CLIPVisionModelWithProjection(UpperCAmelCase_				)
					__UpperCAmelCase   :      List[Any]										=     CLIPImageProcessor(crop_size=32 ,       size=32				)

					__UpperCAmelCase   :      str										=     {
					    'unet': unet,
					    'scheduler': scheduler,
					    'vae': vae,
					    'image_encoder': image_encoder,
					    'feature_extractor': feature_extractor,
					    'safety_checker': None,
					}
					return components





		def 							__A       (     self ,       __UpperCAmelCase ,       __UpperCAmelCase=0				)       ->     Dict:

					'''simple docstring'''
					__UpperCAmelCase   :      Dict										=     floats_tensor((1, 3, 32, 32) ,       rng=random.Random(UpperCAmelCase_				)				)
					__UpperCAmelCase   :      Optional[int]										=     image.cpu().permute(0 ,       2 ,       3 ,       1				)[0]
					__UpperCAmelCase   :      Dict										=     Image.fromarray(np.uinta(UpperCAmelCase_				)				).convert("""RGB"""				).resize((32, 32)				)
					if str(UpperCAmelCase_				).startswith("""mps"""				):
								__UpperCAmelCase   :      str										=     torch.manual_seed(UpperCAmelCase_				)
					else:
								__UpperCAmelCase   :      Optional[int]										=     torch.Generator(device=UpperCAmelCase_				).manual_seed(UpperCAmelCase_				)
					__UpperCAmelCase   :      List[str]										=     {
					    'image': image,
					    'generator': generator,
					    'num_inference_steps': 2,
					    'guidance_scale': 6.0,
					    'output_type': 'numpy',
					}
					return inputs





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

					'''simple docstring'''
					__UpperCAmelCase   :      List[str]										=     'cpu'  # ensure determinism for the device-dependent torch.Generator
					__UpperCAmelCase   :      List[str]										=     self.get_dummy_components()
					__UpperCAmelCase   :      List[Any]										=     StableDiffusionImageVariationPipeline(**UpperCAmelCase_				)
					__UpperCAmelCase   :      Optional[Any]										=     sd_pipe.to(UpperCAmelCase_				)
					sd_pipe.set_progress_bar_config(disable=UpperCAmelCase_				)

					__UpperCAmelCase   :      Optional[int]										=     self.get_dummy_inputs(UpperCAmelCase_				)
					__UpperCAmelCase   :      List[Any]										=     sd_pipe(**UpperCAmelCase_				).images
					__UpperCAmelCase   :      Any										=     image[0, -3:, -3:, -1]

					assert image.shape == (1, 64, 64, 3)
					__UpperCAmelCase   :      Any										=     np.array([0.5239, 0.5723, 0.4796, 0.5049, 0.5550, 0.4685, 0.5329, 0.4891, 0.4921]				)

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





		def 							__A       (     self				)       ->     int:

					'''simple docstring'''
					__UpperCAmelCase   :      Dict										=     'cpu'  # ensure determinism for the device-dependent torch.Generator
					__UpperCAmelCase   :      Any										=     self.get_dummy_components()
					__UpperCAmelCase   :      Any										=     StableDiffusionImageVariationPipeline(**UpperCAmelCase_				)
					__UpperCAmelCase   :      int										=     sd_pipe.to(UpperCAmelCase_				)
					sd_pipe.set_progress_bar_config(disable=UpperCAmelCase_				)

					__UpperCAmelCase   :      Tuple										=     self.get_dummy_inputs(UpperCAmelCase_				)
					__UpperCAmelCase   :      List[str]										=     2 * [inputs['image']]
					__UpperCAmelCase   :      Optional[int]										=     sd_pipe(**UpperCAmelCase_				)

					__UpperCAmelCase   :      List[str]										=     output.images

					__UpperCAmelCase   :      int										=     image[-1, -3:, -3:, -1]

					assert image.shape == (2, 64, 64, 3)
					__UpperCAmelCase   :      Optional[int]										=     np.array([0.6892, 0.5637, 0.5836, 0.5771, 0.6254, 0.6409, 0.5580, 0.5569, 0.5289]				)

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





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

					'''simple docstring'''
					super().test_inference_batch_single_identical(expected_max_diff=3E-3				)



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





		def 							__A       (     self				)       ->     List[Any]:

					'''simple docstring'''
					super().tearDown()
					gc.collect()
					torch.cuda.empty_cache()





		def 							__A       (     self ,       __UpperCAmelCase ,       __UpperCAmelCase="cpu" ,       __UpperCAmelCase=torch.floataa ,       __UpperCAmelCase=0				)       ->     Optional[Any]:

					'''simple docstring'''
					__UpperCAmelCase   :      List[Any]										=     torch.Generator(device=UpperCAmelCase_				).manual_seed(UpperCAmelCase_				)
					__UpperCAmelCase   :      List[str]										=     load_image(
					    """https://huggingface.co/datasets/diffusers/test-arrays/resolve/main"""
					    """/stable_diffusion_imgvar/input_image_vermeer.png"""				)
					__UpperCAmelCase   :      Any										=     np.random.RandomState(UpperCAmelCase_				).standard_normal((1, 4, 64, 64)				)
					__UpperCAmelCase   :      Tuple										=     torch.from_numpy(UpperCAmelCase_				).to(device=UpperCAmelCase_ ,       dtype=UpperCAmelCase_				)
					__UpperCAmelCase   :      str										=     {
					    'image': init_image,
					    'latents': latents,
					    'generator': generator,
					    'num_inference_steps': 3,
					    'guidance_scale': 7.5,
					    'output_type': 'numpy',
					}
					return inputs





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

					'''simple docstring'''
					__UpperCAmelCase   :      Dict										=     StableDiffusionImageVariationPipeline.from_pretrained(
					    """lambdalabs/sd-image-variations-diffusers""" ,       safety_checker=UpperCAmelCase_				)
					__UpperCAmelCase   :      int										=     sd_pipe.to(UpperCAmelCase_				)
					sd_pipe.set_progress_bar_config(disable=UpperCAmelCase_				)

					__UpperCAmelCase   :      Any										=     self.get_inputs(UpperCAmelCase_				)
					__UpperCAmelCase   :      List[str]										=     sd_pipe(**UpperCAmelCase_				).images
					__UpperCAmelCase   :      int										=     image[0, -3:, -3:, -1].flatten()

					assert image.shape == (1, 512, 512, 3)
					__UpperCAmelCase   :      Optional[Any]										=     np.array([0.8_4491, 0.9_0789, 0.7_5708, 0.7_8734, 0.8_3485, 0.7_0099, 0.6_6938, 0.6_8727, 0.6_1379]				)
					assert np.abs(image_slice - expected_slice				).max() < 6E-3





		def 							__A       (     self				)       ->     Optional[Any]:

					'''simple docstring'''
					__UpperCAmelCase   :      str										=     0

					def callback_fn(__UpperCAmelCase ,       __UpperCAmelCase ,       __UpperCAmelCase				) -> None:
								__UpperCAmelCase   :      Optional[int]										=     True
								nonlocal number_of_steps
								number_of_steps += 1
								if step == 1:
											__UpperCAmelCase   :      List[str]										=     latents.detach().cpu().numpy()
											assert latents.shape == (1, 4, 64, 64)
											__UpperCAmelCase   :      Dict										=     latents[0, -3:, -3:, -1]
											__UpperCAmelCase   :      Any										=     np.array(
											    [-0.1621, 0.2837, -0.7979, -0.1221, -1.3057, 0.7681, -2.1191, 0.0464, 1.6309]				)

											assert np.abs(latents_slice.flatten() - expected_slice				).max() < 5E-2
								elif step == 2:
											__UpperCAmelCase   :      List[Any]										=     latents.detach().cpu().numpy()
											assert latents.shape == (1, 4, 64, 64)
											__UpperCAmelCase   :      Union[str, Any]										=     latents[0, -3:, -3:, -1]
											__UpperCAmelCase   :      Optional[int]										=     np.array([0.6299, 1.7500, 1.1992, -2.1582, -1.8994, 0.7334, -0.7090, 1.0137, 1.5273]				)

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

					__UpperCAmelCase   :      Dict										=     False

					__UpperCAmelCase   :      List[Any]										=     StableDiffusionImageVariationPipeline.from_pretrained(
					    """fusing/sd-image-variations-diffusers""" ,       safety_checker=UpperCAmelCase_ ,       torch_dtype=torch.floataa ,       )
					pipe.to(UpperCAmelCase_				)
					pipe.set_progress_bar_config(disable=UpperCAmelCase_				)
					pipe.enable_attention_slicing()

					__UpperCAmelCase   :      Optional[Any]										=     self.get_inputs(UpperCAmelCase_ ,       dtype=torch.floataa				)
					pipe(**UpperCAmelCase_ ,       callback=UpperCAmelCase_ ,       callback_steps=1				)
					assert callback_fn.has_been_called
					assert number_of_steps == inputs["num_inference_steps"]





		def 							__A       (     self				)       ->     int:

					'''simple docstring'''
					torch.cuda.empty_cache()
					torch.cuda.reset_max_memory_allocated()
					torch.cuda.reset_peak_memory_stats()

					__UpperCAmelCase   :      Optional[int]										=     'fusing/sd-image-variations-diffusers'
					__UpperCAmelCase   :      int										=     StableDiffusionImageVariationPipeline.from_pretrained(
					    UpperCAmelCase_ ,       safety_checker=UpperCAmelCase_ ,       torch_dtype=torch.floataa				)
					__UpperCAmelCase   :      int										=     pipe.to(UpperCAmelCase_				)
					pipe.set_progress_bar_config(disable=UpperCAmelCase_				)
					pipe.enable_attention_slicing(1				)
					pipe.enable_sequential_cpu_offload()

					__UpperCAmelCase   :      str										=     self.get_inputs(UpperCAmelCase_ ,       dtype=torch.floataa				)
					__UpperCAmelCase   :      List[str]										=     pipe(**UpperCAmelCase_				)

					__UpperCAmelCase   :      int										=     torch.cuda.max_memory_allocated()
					# make sure that less than 2.6 GB is allocated
					assert mem_bytes < 2.6 * 10**9



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





		def 							__A       (     self				)       ->     int:

					'''simple docstring'''
					super().tearDown()
					gc.collect()
					torch.cuda.empty_cache()





		def 							__A       (     self ,       __UpperCAmelCase ,       __UpperCAmelCase="cpu" ,       __UpperCAmelCase=torch.floataa ,       __UpperCAmelCase=0				)       ->     int:

					'''simple docstring'''
					__UpperCAmelCase   :      int										=     torch.Generator(device=UpperCAmelCase_				).manual_seed(UpperCAmelCase_				)
					__UpperCAmelCase   :      Optional[Any]										=     load_image(
					    """https://huggingface.co/datasets/diffusers/test-arrays/resolve/main"""
					    """/stable_diffusion_imgvar/input_image_vermeer.png"""				)
					__UpperCAmelCase   :      Tuple										=     np.random.RandomState(UpperCAmelCase_				).standard_normal((1, 4, 64, 64)				)
					__UpperCAmelCase   :      int										=     torch.from_numpy(UpperCAmelCase_				).to(device=UpperCAmelCase_ ,       dtype=UpperCAmelCase_				)
					__UpperCAmelCase   :      List[str]										=     {
					    'image': init_image,
					    'latents': latents,
					    'generator': generator,
					    'num_inference_steps': 50,
					    'guidance_scale': 7.5,
					    'output_type': 'numpy',
					}
					return inputs





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

					'''simple docstring'''
					__UpperCAmelCase   :      Union[str, Any]										=     StableDiffusionImageVariationPipeline.from_pretrained("""fusing/sd-image-variations-diffusers"""				)
					sd_pipe.to(UpperCAmelCase_				)
					sd_pipe.set_progress_bar_config(disable=UpperCAmelCase_				)

					__UpperCAmelCase   :      Optional[Any]										=     self.get_inputs(UpperCAmelCase_				)
					__UpperCAmelCase   :      int										=     sd_pipe(**UpperCAmelCase_				).images[0]

					__UpperCAmelCase   :      int										=     load_numpy(
					    """https://huggingface.co/datasets/diffusers/test-arrays/resolve/main"""
					    """/stable_diffusion_imgvar/lambdalabs_variations_pndm.npy"""				)
					__UpperCAmelCase   :      Optional[Any]										=     np.abs(expected_image - image				).max()
					assert max_diff < 1E-3





		def 							__A       (     self				)       ->     Dict:

					'''simple docstring'''
					__UpperCAmelCase   :      Optional[Any]										=     StableDiffusionImageVariationPipeline.from_pretrained("""fusing/sd-image-variations-diffusers"""				)
					__UpperCAmelCase   :      str										=     DPMSolverMultistepScheduler.from_config(sd_pipe.scheduler.config				)
					sd_pipe.to(UpperCAmelCase_				)
					sd_pipe.set_progress_bar_config(disable=UpperCAmelCase_				)

					__UpperCAmelCase   :      Dict										=     self.get_inputs(UpperCAmelCase_				)
					__UpperCAmelCase   :      Optional[Any]										=     25
					__UpperCAmelCase   :      Optional[Any]										=     sd_pipe(**UpperCAmelCase_				).images[0]

					__UpperCAmelCase   :      Tuple										=     load_numpy(
					    """https://huggingface.co/datasets/diffusers/test-arrays/resolve/main"""
					    """/stable_diffusion_imgvar/lambdalabs_variations_dpm_multi.npy"""				)
					__UpperCAmelCase   :      str										=     np.abs(expected_image - image				).max()
					assert max_diff < 1E-3









'''simple docstring'''
from collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


_UpperCamelCase											=			logging.get_logger(__name__)

_UpperCamelCase											=			{
    '''nghuyong/ernie-1.0-base-zh''': '''https://huggingface.co/nghuyong/ernie-1.0-base-zh/resolve/main/config.json''',
    '''nghuyong/ernie-2.0-base-en''': '''https://huggingface.co/nghuyong/ernie-2.0-base-en/resolve/main/config.json''',
    '''nghuyong/ernie-2.0-large-en''': '''https://huggingface.co/nghuyong/ernie-2.0-large-en/resolve/main/config.json''',
    '''nghuyong/ernie-3.0-base-zh''': '''https://huggingface.co/nghuyong/ernie-3.0-base-zh/resolve/main/config.json''',
    '''nghuyong/ernie-3.0-medium-zh''': '''https://huggingface.co/nghuyong/ernie-3.0-medium-zh/resolve/main/config.json''',
    '''nghuyong/ernie-3.0-mini-zh''': '''https://huggingface.co/nghuyong/ernie-3.0-mini-zh/resolve/main/config.json''',
    '''nghuyong/ernie-3.0-micro-zh''': '''https://huggingface.co/nghuyong/ernie-3.0-micro-zh/resolve/main/config.json''',
    '''nghuyong/ernie-3.0-nano-zh''': '''https://huggingface.co/nghuyong/ernie-3.0-nano-zh/resolve/main/config.json''',
    '''nghuyong/ernie-gram-zh''': '''https://huggingface.co/nghuyong/ernie-gram-zh/resolve/main/config.json''',
    '''nghuyong/ernie-health-zh''': '''https://huggingface.co/nghuyong/ernie-health-zh/resolve/main/config.json''',
}



class _A  (    __SCREAMING_SNAKE_CASE    ):
		_SCREAMING_SNAKE_CASE					:    Optional[int]						      =			"ernie"





		def __init__(     self ,       __UpperCAmelCase=30_522 ,       __UpperCAmelCase=768 ,       __UpperCAmelCase=12 ,       __UpperCAmelCase=12 ,       __UpperCAmelCase=3_072 ,       __UpperCAmelCase="gelu" ,       __UpperCAmelCase=0.1 ,       __UpperCAmelCase=0.1 ,       __UpperCAmelCase=512 ,       __UpperCAmelCase=2 ,       __UpperCAmelCase=3 ,       __UpperCAmelCase=False ,       __UpperCAmelCase=0.02 ,       __UpperCAmelCase=1E-12 ,       __UpperCAmelCase=0 ,       __UpperCAmelCase="absolute" ,       __UpperCAmelCase=True ,       __UpperCAmelCase=None ,       **__UpperCAmelCase ,       )       ->     List[Any]:

					'''simple docstring'''
					super().__init__(pad_token_id=__UpperCAmelCase ,       **__UpperCAmelCase				)

					__UpperCAmelCase   :      List[str]										=     vocab_size
					__UpperCAmelCase   :      Optional[int]										=     hidden_size
					__UpperCAmelCase   :      Optional[Any]										=     num_hidden_layers
					__UpperCAmelCase   :      int										=     num_attention_heads
					__UpperCAmelCase   :      int										=     hidden_act
					__UpperCAmelCase   :      List[Any]										=     intermediate_size
					__UpperCAmelCase   :      Optional[Any]										=     hidden_dropout_prob
					__UpperCAmelCase   :      Dict										=     attention_probs_dropout_prob
					__UpperCAmelCase   :      int										=     max_position_embeddings
					__UpperCAmelCase   :      str										=     type_vocab_size
					__UpperCAmelCase   :      int										=     task_type_vocab_size
					__UpperCAmelCase   :      List[str]										=     use_task_id
					__UpperCAmelCase   :      Union[str, Any]										=     initializer_range
					__UpperCAmelCase   :      Tuple										=     layer_norm_eps
					__UpperCAmelCase   :      List[str]										=     position_embedding_type
					__UpperCAmelCase   :      Union[str, Any]										=     use_cache
					__UpperCAmelCase   :      Optional[int]										=     classifier_dropout



class _A  (    __SCREAMING_SNAKE_CASE    ):





		@property
		def 							__A       (     self				)       ->     Mapping[str, Mapping[int, str]]:

					'''simple docstring'''
					if self.task == "multiple-choice":
								__UpperCAmelCase   :      Union[str, Any]										=     {0: """batch""", 1: """choice""", 2: """sequence"""}
					else:
								__UpperCAmelCase   :      Tuple										=     {0: """batch""", 1: """sequence"""}
					return OrderedDict(
					    [
					        ("""input_ids""", dynamic_axis),
					        ("""attention_mask""", dynamic_axis),
					        ("""token_type_ids""", dynamic_axis),
					        ("""task_type_ids""", dynamic_axis),
					    ]				)




