


import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModelWithProjection, CLIPTokenizer

from diffusers import HeunDiscreteScheduler, PriorTransformer, ShapEPipeline
from diffusers.pipelines.shap_e import ShapERenderer
from diffusers.utils import load_numpy, slow
from diffusers.utils.testing_utils import require_torch_gpu, torch_device

from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference






class     UpperCamelCase(  SCREAMING_SNAKE_CASE__     ,							unittest.TestCase   ):
					snake_case_		:				Optional[Any]										=	ShapEPipeline
					snake_case_		:				List[Any]										=	["""prompt"""]
					snake_case_		:				List[str]										=	["""prompt"""]
					snake_case_		:				Optional[Any]										=	[
					    """num_images_per_prompt""",
					    """num_inference_steps""",
					    """generator""",
					    """latents""",
					    """guidance_scale""",
					    """frame_size""",
					    """output_type""",
					    """return_dict""",
					]
					snake_case_		:				Optional[int]										=	False


					@property
					def 							SCREAMING_SNAKE_CASE_  (						self	:		Dict						)			->	int:




											'''simple docstring'''

											return 3_2


					@property
					def 							SCREAMING_SNAKE_CASE_  (						self	:		List[str]						)			->	Union[str, Any]:




											'''simple docstring'''

											return 3_2


					@property
					def 							SCREAMING_SNAKE_CASE_  (						self	:		Optional[int]						)			->	Any:




											'''simple docstring'''

											return self.time_input_dim * 4


					@property
					def 							SCREAMING_SNAKE_CASE_  (						self	:		List[Any]						)			->	Optional[int]:




											'''simple docstring'''

											return 8


					@property
					def 							SCREAMING_SNAKE_CASE_  (						self	:		int						)			->	List[str]:




											'''simple docstring'''

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


					@property
					def 							SCREAMING_SNAKE_CASE_  (						self	:		Tuple						)			->	List[str]:




											'''simple docstring'''

											torch.manual_seed(0						)
											__snake_case				=		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=3_7     ,    layer_norm_eps=1e-0_5     ,    num_attention_heads=4     ,    num_hidden_layers=5     ,    pad_token_id=1     ,    vocab_size=1_0_0_0     ,    )
											return CLIPTextModelWithProjection(_lowercase						)


					@property
					def 							SCREAMING_SNAKE_CASE_  (						self	:		Optional[Any]						)			->	List[Any]:




											'''simple docstring'''

											torch.manual_seed(0						)

											__snake_case				=		{
											    """num_attention_heads""": 2,
											    """attention_head_dim""": 1_6,
											    """embedding_dim""": self.time_input_dim,
											    """num_embeddings""": 3_2,
											    """embedding_proj_dim""": self.text_embedder_hidden_size,
											    """time_embed_dim""": self.time_embed_dim,
											    """num_layers""": 1,
											    """clip_embed_dim""": self.time_input_dim * 2,
											    """additional_embeddings""": 0,
											    """time_embed_act_fn""": """gelu""",
											    """norm_in_type""": """layer""",
											    """encoder_hid_proj_type""": None,
											    """added_emb_type""": None,
											}

											__snake_case				=		PriorTransformer(**_lowercase						)
											return model


					@property
					def 							SCREAMING_SNAKE_CASE_  (						self	:		Dict						)			->	Any:




											'''simple docstring'''

											torch.manual_seed(0						)

											__snake_case				=		{
											    """param_shapes""": (
											        (self.renderer_dim, 9_3),
											        (self.renderer_dim, 8),
											        (self.renderer_dim, 8),
											        (self.renderer_dim, 8),
											    ),
											    """d_latent""": self.time_input_dim,
											    """d_hidden""": self.renderer_dim,
											    """n_output""": 1_2,
											    """background""": (
											        0.1,
											        0.1,
											        0.1,
											    ),
											}
											__snake_case				=		ShapERenderer(**_lowercase						)
											return model


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




											'''simple docstring'''

											__snake_case				=		self.dummy_prior
											__snake_case				=		self.dummy_text_encoder
											__snake_case				=		self.dummy_tokenizer
											__snake_case				=		self.dummy_renderer

											__snake_case				=		HeunDiscreteScheduler(
											    beta_schedule="exp"     ,    num_train_timesteps=1_0_2_4     ,    prediction_type="sample"     ,    use_karras_sigmas=_lowercase     ,    clip_sample=_lowercase     ,    clip_sample_range=1.0     ,    )
											__snake_case				=		{
											    """prior""": prior,
											    """text_encoder""": text_encoder,
											    """tokenizer""": tokenizer,
											    """renderer""": renderer,
											    """scheduler""": scheduler,
											}

											return components


					def 							SCREAMING_SNAKE_CASE_  (						self	:		Union[str, Any]     ,    SCREAMING_SNAKE_CASE	:		Optional[Any]     ,    SCREAMING_SNAKE_CASE	:		str=0						)			->	int:




											'''simple docstring'''

											if str(_lowercase						).startswith("mps"						):
																	__snake_case				=		torch.manual_seed(_lowercase						)
											else:
																	__snake_case				=		torch.Generator(device=_lowercase						).manual_seed(_lowercase						)
											__snake_case				=		{
											    """prompt""": """horse""",
											    """generator""": generator,
											    """num_inference_steps""": 1,
											    """frame_size""": 3_2,
											    """output_type""": """np""",
											}
											return inputs


					def 							SCREAMING_SNAKE_CASE_  (						self	:		Dict						)			->	int:




											'''simple docstring'''

											__snake_case				=		"""cpu"""

											__snake_case				=		self.get_dummy_components()

											__snake_case				=		self.pipeline_class(**_lowercase						)
											__snake_case				=		pipe.to(_lowercase						)

											pipe.set_progress_bar_config(disable=_lowercase						)

											__snake_case				=		pipe(**self.get_dummy_inputs(_lowercase						)						)
											__snake_case				=		output.images[0]
											__snake_case				=		image[0, -3:, -3:, -1]

											assert image.shape == (2_0, 3_2, 3_2, 3)

											__snake_case				=		np.array(
											    [
											        0.00039216,
											        0.00039216,
											        0.00039216,
											        0.00039216,
											        0.00039216,
											        0.00039216,
											        0.00039216,
											        0.00039216,
											        0.00039216,
											    ]						)

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


					def 							SCREAMING_SNAKE_CASE_  (						self	:		int						)			->	List[str]:




											'''simple docstring'''

											self._test_inference_batch_consistent(batch_sizes=[1, 2]						)


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




											'''simple docstring'''

											__snake_case				=		torch_device == """cpu"""
											__snake_case				=		True

											self._test_inference_batch_single_identical(
											    batch_size=2     ,    test_max_difference=_lowercase     ,    relax_max_difference=_lowercase     ,    )






					def 							SCREAMING_SNAKE_CASE_  (						self	:		int						)			->	Dict:




											'''simple docstring'''

											__snake_case				=		self.get_dummy_components()
											__snake_case				=		self.pipeline_class(**_lowercase						)
											__snake_case				=		pipe.to(_lowercase						)
											pipe.set_progress_bar_config(disable=_lowercase						)

											__snake_case				=		1
											__snake_case				=		2

											__snake_case				=		self.get_dummy_inputs(_lowercase						)

											for key in inputs.keys():
																	if key in self.batch_params:
																							__snake_case				=		batch_size * [inputs[key]]

											__snake_case				=		pipe(**_lowercase     ,    num_images_per_prompt=_lowercase						)[0]

											assert images.shape[0] == batch_size * num_images_per_prompt






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


					def 							SCREAMING_SNAKE_CASE_  (						self	:		int						)			->	Optional[Any]:




											'''simple docstring'''

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






					def 							SCREAMING_SNAKE_CASE_  (						self	:		Optional[int]						)			->	Dict:




											'''simple docstring'''

											__snake_case				=		load_numpy(
											    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
											    "/shap_e/test_shap_e_np_out.npy"						)
											__snake_case				=		ShapEPipeline.from_pretrained("openai/shap-e"						)
											__snake_case				=		pipe.to(_lowercase						)
											pipe.set_progress_bar_config(disable=_lowercase						)

											__snake_case				=		torch.Generator(device=_lowercase						).manual_seed(0						)

											__snake_case				=		pipe(
											    "a shark"     ,    generator=_lowercase     ,    guidance_scale=15.0     ,    num_inference_steps=6_4     ,    frame_size=6_4     ,    output_type="np"     ,    ).images[0]

											assert images.shape == (2_0, 6_4, 6_4, 3)

											assert_mean_pixel_difference(_lowercase     ,    _lowercase						)









import random
import unittest

import numpy as np
import torch

from diffusers import (
    DPMSolverMultistepScheduler,
    EulerAncestralDiscreteScheduler,
    EulerDiscreteScheduler,
    LMSDiscreteScheduler,
    OnnxStableDiffusionUpscalePipeline,
    PNDMScheduler,
)
from diffusers.utils import floats_tensor
from diffusers.utils.testing_utils import (
    is_onnx_available,
    load_image,
    nightly,
    require_onnxruntime,
    require_torch_gpu,
)

from ..test_pipelines_onnx_common import OnnxPipelineTesterMixin


if is_onnx_available():
	import onnxruntime as ort






class     UpperCamelCase(  _a     ,							unittest.TestCase   ):
					# TODO: is there an appropriate internal test set?
					snake_case_		:				Dict										=	"""ssube/stable-diffusion-x4-upscaler-onnx"""


					def 							SCREAMING_SNAKE_CASE_  (						self	:		Dict     ,    SCREAMING_SNAKE_CASE	:		List[str]=0						)			->	Union[str, Any]:




											'''simple docstring'''

											__snake_case				=		floats_tensor((1, 3, 1_2_8, 1_2_8)     ,    rng=random.Random(SCREAMING_SNAKE_CASE						)						)
											__snake_case				=		torch.manual_seed(SCREAMING_SNAKE_CASE						)
											__snake_case				=		{
											    "prompt": "A painting of a squirrel eating a burger",
											    "image": image,
											    "generator": generator,
											    "num_inference_steps": 3,
											    "guidance_scale": 7.5,
											    "output_type": "numpy",
											}
											return inputs


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




											'''simple docstring'''

											__snake_case				=		OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint     ,    provider="CPUExecutionProvider"						)
											pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE						)

											__snake_case				=		self.get_dummy_inputs()
											__snake_case				=		pipe(**SCREAMING_SNAKE_CASE						).images
											__snake_case				=		image[0, -3:, -3:, -1].flatten()

											# started as 128, should now be 512
											assert image.shape == (1, 5_1_2, 5_1_2, 3)
											__snake_case				=		np.array(
											    [0.6974782, 0.68902093, 0.70135885, 0.7583618, 0.7804545, 0.7854912, 0.78667426, 0.78743863, 0.78070223]						)
											assert np.abs(image_slice - expected_slice						).max() < 1e-1


					def 							SCREAMING_SNAKE_CASE_  (						self	:		List[str]						)			->	Optional[Any]:




											'''simple docstring'''

											__snake_case				=		OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint     ,    provider="CPUExecutionProvider"						)
											__snake_case				=		PNDMScheduler.from_config(pipe.scheduler.config     ,    skip_prk_steps=SCREAMING_SNAKE_CASE						)
											pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE						)

											__snake_case				=		self.get_dummy_inputs()
											__snake_case				=		pipe(**SCREAMING_SNAKE_CASE						).images
											__snake_case				=		image[0, -3:, -3:, -1]

											assert image.shape == (1, 5_1_2, 5_1_2, 3)
											__snake_case				=		np.array(
											    [0.6898892, 0.59240556, 0.52499527, 0.58866215, 0.52258235, 0.52572715, 0.62414473, 0.6174387, 0.6214964]						)
											assert np.abs(image_slice.flatten() - expected_slice						).max() < 1e-1


					def 							SCREAMING_SNAKE_CASE_  (						self	:		Any						)			->	Optional[int]:




											'''simple docstring'''

											__snake_case				=		OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint     ,    provider="CPUExecutionProvider"						)
											__snake_case				=		DPMSolverMultistepScheduler.from_config(pipe.scheduler.config						)
											pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE						)

											__snake_case				=		self.get_dummy_inputs()
											__snake_case				=		pipe(**SCREAMING_SNAKE_CASE						).images
											__snake_case				=		image[0, -3:, -3:, -1]

											assert image.shape == (1, 5_1_2, 5_1_2, 3)
											__snake_case				=		np.array(
											    [0.7659278, 0.76437664, 0.75579107, 0.7691116, 0.77666986, 0.7727672, 0.7758664, 0.7812226, 0.76942515]						)

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


					def 							SCREAMING_SNAKE_CASE_  (						self	:		Dict						)			->	Dict:




											'''simple docstring'''

											__snake_case				=		OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint     ,    provider="CPUExecutionProvider"						)
											__snake_case				=		EulerDiscreteScheduler.from_config(pipe.scheduler.config						)
											pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE						)

											__snake_case				=		self.get_dummy_inputs()
											__snake_case				=		pipe(**SCREAMING_SNAKE_CASE						).images
											__snake_case				=		image[0, -3:, -3:, -1]

											assert image.shape == (1, 5_1_2, 5_1_2, 3)
											__snake_case				=		np.array(
											    [0.6974782, 0.68902093, 0.70135885, 0.7583618, 0.7804545, 0.7854912, 0.78667426, 0.78743863, 0.78070223]						)
											assert np.abs(image_slice.flatten() - expected_slice						).max() < 1e-1






					def 							SCREAMING_SNAKE_CASE_  (						self	:		List[Any]						)			->	List[str]:




											'''simple docstring'''

											__snake_case				=		OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint     ,    provider="CPUExecutionProvider"						)
											__snake_case				=		EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config						)
											pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE						)

											__snake_case				=		self.get_dummy_inputs()
											__snake_case				=		pipe(**SCREAMING_SNAKE_CASE						).images
											__snake_case				=		image[0, -3:, -3:, -1]

											assert image.shape == (1, 5_1_2, 5_1_2, 3)
											__snake_case				=		np.array(
											    [0.77424496, 0.773601, 0.7645288, 0.7769598, 0.7772739, 0.7738688, 0.78187233, 0.77879584, 0.767043]						)

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






@nightly
@require_onnxruntime
@require_torch_gpu
class     UpperCamelCase(  unittest.TestCase   ):


					@property
					def 							SCREAMING_SNAKE_CASE_  (						self	:		str						)			->	List[Any]:




											'''simple docstring'''

											return (
											    "CUDAExecutionProvider",
											    {
											        "gpu_mem_limit": "15000000000",  # 15GB
											        "arena_extend_strategy": "kSameAsRequested",
											    },
											)


					@property
					def 							SCREAMING_SNAKE_CASE_  (						self	:		Tuple						)			->	Tuple:




											'''simple docstring'''

											__snake_case				=		ort.SessionOptions()
											__snake_case				=		False
											return options


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




											'''simple docstring'''

											__snake_case				=		load_image(
											    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
											    "/img2img/sketch-mountains-input.jpg"						)
											__snake_case				=		init_image.resize((1_2_8, 1_2_8)						)
											# using the PNDM scheduler by default
											__snake_case				=		OnnxStableDiffusionUpscalePipeline.from_pretrained(
											    "ssube/stable-diffusion-x4-upscaler-onnx"     ,    provider=self.gpu_provider     ,    sess_options=self.gpu_options     ,    )
											pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE						)

											__snake_case				=		"A fantasy landscape, trending on artstation"

											__snake_case				=		torch.manual_seed(0						)
											__snake_case				=		pipe(
											    prompt=SCREAMING_SNAKE_CASE     ,    image=SCREAMING_SNAKE_CASE     ,    guidance_scale=7.5     ,    num_inference_steps=1_0     ,    generator=SCREAMING_SNAKE_CASE     ,    output_type="np"     ,    )
											__snake_case				=		output.images
											__snake_case				=		images[0, 2_5_5:2_5_8, 3_8_3:3_8_6, -1]

											assert images.shape == (1, 5_1_2, 5_1_2, 3)
											__snake_case				=		np.array([0.4883, 0.4947, 0.4980, 0.4975, 0.4982, 0.4980, 0.5000, 0.5006, 0.4972]						)
											# TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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






					def 							SCREAMING_SNAKE_CASE_  (						self	:		Tuple						)			->	Optional[int]:




											'''simple docstring'''

											__snake_case				=		load_image(
											    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
											    "/img2img/sketch-mountains-input.jpg"						)
											__snake_case				=		init_image.resize((1_2_8, 1_2_8)						)
											__snake_case				=		LMSDiscreteScheduler.from_pretrained(
											    "ssube/stable-diffusion-x4-upscaler-onnx"     ,    subfolder="scheduler"						)
											__snake_case				=		OnnxStableDiffusionUpscalePipeline.from_pretrained(
											    "ssube/stable-diffusion-x4-upscaler-onnx"     ,    scheduler=SCREAMING_SNAKE_CASE     ,    provider=self.gpu_provider     ,    sess_options=self.gpu_options     ,    )
											pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE						)

											__snake_case				=		"A fantasy landscape, trending on artstation"

											__snake_case				=		torch.manual_seed(0						)
											__snake_case				=		pipe(
											    prompt=SCREAMING_SNAKE_CASE     ,    image=SCREAMING_SNAKE_CASE     ,    guidance_scale=7.5     ,    num_inference_steps=2_0     ,    generator=SCREAMING_SNAKE_CASE     ,    output_type="np"     ,    )
											__snake_case				=		output.images
											__snake_case				=		images[0, 2_5_5:2_5_8, 3_8_3:3_8_6, -1]

											assert images.shape == (1, 5_1_2, 5_1_2, 3)
											__snake_case				=		np.array(
											    [0.50173753, 0.50223356, 0.502039, 0.50233036, 0.5023725, 0.5022601, 0.5018758, 0.50234085, 0.50241566]						)
											# TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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





