





import random
import unittest

import numpy as np

from diffusers import (
    DPMSolverMultistepScheduler,
    EulerAncestralDiscreteScheduler,
    EulerDiscreteScheduler,
    LMSDiscreteScheduler,
    OnnxStableDiffusionImgaImgPipeline,
    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 						lowerCamelCase       (__lowerCamelCase							,     unittest.TestCase      ):







							"""simple docstring"""




							UpperCAmelCase_      						=   "hf-internal-testing/tiny-random-OnnxStableDiffusionPipeline"






							def  A_     ( self     :					int,			_UpperCAmelCase     :					Optional[int]=0						)				->			List[str]:





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__       :      List[str]     =						floats_tensor((1, 3, 1_2_8, 1_2_8),			rng=random.Random(_UpperCAmelCase						)						)
														SCREAMING_SNAKE_CASE__       :      str     =						np.random.RandomState(_UpperCAmelCase						)
														SCREAMING_SNAKE_CASE__       :      Any     =						{
														    "prompt": "A painting of a squirrel eating a burger",
														    "image": image,
														    "generator": generator,
														    "num_inference_steps": 3,
														    "strength": 0.75,
														    "guidance_scale": 7.5,
														    "output_type": "numpy",
														}
														return inputs






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





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__       :      Union[str, Any]     =						OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint,			provider="CPUExecutionProvider"						)
														pipe.set_progress_bar_config(disable=_UpperCAmelCase						)

														SCREAMING_SNAKE_CASE__       :      Tuple     =						self.get_dummy_inputs()
														SCREAMING_SNAKE_CASE__       :      Optional[int]     =						pipe(**_UpperCAmelCase						).images
														SCREAMING_SNAKE_CASE__       :      Any     =						image[0, -3:, -3:, -1].flatten()

														assert image.shape == (1, 1_2_8, 1_2_8, 3)
														SCREAMING_SNAKE_CASE__       :      Union[str, Any]     =						np.array([0.69643, 0.58484, 0.50314, 0.58760, 0.55368, 0.59643, 0.51529, 0.41217, 0.49087]						)
														assert np.abs(image_slice - expected_slice						).max() < 1E-1






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





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__       :      int     =						OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint,			provider="CPUExecutionProvider"						)
														SCREAMING_SNAKE_CASE__       :      List[str]     =						PNDMScheduler.from_config(pipe.scheduler.config,			skip_prk_steps=_UpperCAmelCase						)
														pipe.set_progress_bar_config(disable=_UpperCAmelCase						)

														SCREAMING_SNAKE_CASE__       :      List[Any]     =						self.get_dummy_inputs()
														SCREAMING_SNAKE_CASE__       :      Dict     =						pipe(**_UpperCAmelCase						).images
														SCREAMING_SNAKE_CASE__       :      Any     =						image[0, -3:, -3:, -1]

														assert image.shape == (1, 1_2_8, 1_2_8, 3)
														SCREAMING_SNAKE_CASE__       :      Any     =						np.array([0.61737, 0.54642, 0.53183, 0.54465, 0.52742, 0.60525, 0.49969, 0.40655, 0.48154]						)

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






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





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__       :      str     =						OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint,			provider="CPUExecutionProvider"						)
														SCREAMING_SNAKE_CASE__       :      Dict     =						LMSDiscreteScheduler.from_config(pipe.scheduler.config						)
														pipe.set_progress_bar_config(disable=_UpperCAmelCase						)

														# warmup pass to apply optimizations
														SCREAMING_SNAKE_CASE__       :      Tuple     =						pipe(**self.get_dummy_inputs()						)

														SCREAMING_SNAKE_CASE__       :      int     =						self.get_dummy_inputs()
														SCREAMING_SNAKE_CASE__       :      Tuple     =						pipe(**_UpperCAmelCase						).images
														SCREAMING_SNAKE_CASE__       :      List[Any]     =						image[0, -3:, -3:, -1]

														assert image.shape == (1, 1_2_8, 1_2_8, 3)
														SCREAMING_SNAKE_CASE__       :      Optional[Any]     =						np.array([0.52761, 0.59977, 0.49033, 0.49619, 0.54282, 0.50311, 0.47600, 0.40918, 0.45203]						)

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






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





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__       :      Dict     =						OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint,			provider="CPUExecutionProvider"						)
														SCREAMING_SNAKE_CASE__       :      Union[str, Any]     =						EulerDiscreteScheduler.from_config(pipe.scheduler.config						)
														pipe.set_progress_bar_config(disable=_UpperCAmelCase						)

														SCREAMING_SNAKE_CASE__       :      Union[str, Any]     =						self.get_dummy_inputs()
														SCREAMING_SNAKE_CASE__       :      Optional[int]     =						pipe(**_UpperCAmelCase						).images
														SCREAMING_SNAKE_CASE__       :      Union[str, Any]     =						image[0, -3:, -3:, -1]

														assert image.shape == (1, 1_2_8, 1_2_8, 3)
														SCREAMING_SNAKE_CASE__       :      str     =						np.array([0.52911, 0.60004, 0.49229, 0.49805, 0.54502, 0.50680, 0.47777, 0.41028, 0.45304]						)

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






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





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__       :      Optional[int]     =						OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint,			provider="CPUExecutionProvider"						)
														SCREAMING_SNAKE_CASE__       :      int     =						EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config						)
														pipe.set_progress_bar_config(disable=_UpperCAmelCase						)

														SCREAMING_SNAKE_CASE__       :      List[str]     =						self.get_dummy_inputs()
														SCREAMING_SNAKE_CASE__       :      Any     =						pipe(**_UpperCAmelCase						).images
														SCREAMING_SNAKE_CASE__       :      str     =						image[0, -3:, -3:, -1]

														assert image.shape == (1, 1_2_8, 1_2_8, 3)
														SCREAMING_SNAKE_CASE__       :      Dict     =						np.array([0.52911, 0.60004, 0.49229, 0.49805, 0.54502, 0.50680, 0.47777, 0.41028, 0.45304]						)

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






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





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__       :      Union[str, Any]     =						OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint,			provider="CPUExecutionProvider"						)
														SCREAMING_SNAKE_CASE__       :      Optional[int]     =						DPMSolverMultistepScheduler.from_config(pipe.scheduler.config						)
														pipe.set_progress_bar_config(disable=_UpperCAmelCase						)

														SCREAMING_SNAKE_CASE__       :      Optional[Any]     =						self.get_dummy_inputs()
														SCREAMING_SNAKE_CASE__       :      str     =						pipe(**_UpperCAmelCase						).images
														SCREAMING_SNAKE_CASE__       :      int     =						image[0, -3:, -3:, -1]

														assert image.shape == (1, 1_2_8, 1_2_8, 3)
														SCREAMING_SNAKE_CASE__       :      Optional[int]     =						np.array([0.65331, 0.58277, 0.48204, 0.56059, 0.53665, 0.56235, 0.50969, 0.40009, 0.46552]						)

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


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







							"""simple docstring"""




							@property
							def  A_     ( self     :					Tuple						)				->			str:





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






							@property
							def  A_     ( self     :					int						)				->			str:





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__       :      Optional[Any]     =						ort.SessionOptions()
														SCREAMING_SNAKE_CASE__       :      Optional[int]     =						False
														return options






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





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__       :      Union[str, Any]     =						load_image(
														    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
														    "/img2img/sketch-mountains-input.jpg"						)
														SCREAMING_SNAKE_CASE__       :      List[Any]     =						init_image.resize((7_6_8, 5_1_2)						)
														# using the PNDM scheduler by default
														SCREAMING_SNAKE_CASE__       :      Optional[Any]     =						OnnxStableDiffusionImgaImgPipeline.from_pretrained(
														    "CompVis/stable-diffusion-v1-4",			revision="onnx",			safety_checker=_UpperCAmelCase,			feature_extractor=_UpperCAmelCase,			provider=self.gpu_provider,			sess_options=self.gpu_options,			)
														pipe.set_progress_bar_config(disable=_UpperCAmelCase						)

														SCREAMING_SNAKE_CASE__       :      Any     =						"A fantasy landscape, trending on artstation"

														SCREAMING_SNAKE_CASE__       :      List[Any]     =						np.random.RandomState(0						)
														SCREAMING_SNAKE_CASE__       :      Optional[int]     =						pipe(
														    prompt=_UpperCAmelCase,			image=_UpperCAmelCase,			strength=0.75,			guidance_scale=7.5,			num_inference_steps=1_0,			generator=_UpperCAmelCase,			output_type="np",			)
														SCREAMING_SNAKE_CASE__       :      List[Any]     =						output.images
														SCREAMING_SNAKE_CASE__       :      List[str]     =						images[0, 2_5_5:2_5_8, 3_8_3:3_8_6, -1]

														assert images.shape == (1, 5_1_2, 7_6_8, 3)
														SCREAMING_SNAKE_CASE__       :      List[Any]     =						np.array([0.4909, 0.5059, 0.5372, 0.4623, 0.4876, 0.5049, 0.4820, 0.4956, 0.5019]						)
														# TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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






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





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__       :      Any     =						load_image(
														    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
														    "/img2img/sketch-mountains-input.jpg"						)
														SCREAMING_SNAKE_CASE__       :      str     =						init_image.resize((7_6_8, 5_1_2)						)
														SCREAMING_SNAKE_CASE__       :      Any     =						LMSDiscreteScheduler.from_pretrained(
														    "runwayml/stable-diffusion-v1-5",			subfolder="scheduler",			revision="onnx"						)
														SCREAMING_SNAKE_CASE__       :      Tuple     =						OnnxStableDiffusionImgaImgPipeline.from_pretrained(
														    "runwayml/stable-diffusion-v1-5",			revision="onnx",			scheduler=_UpperCAmelCase,			safety_checker=_UpperCAmelCase,			feature_extractor=_UpperCAmelCase,			provider=self.gpu_provider,			sess_options=self.gpu_options,			)
														pipe.set_progress_bar_config(disable=_UpperCAmelCase						)

														SCREAMING_SNAKE_CASE__       :      Optional[Any]     =						"A fantasy landscape, trending on artstation"

														SCREAMING_SNAKE_CASE__       :      Any     =						np.random.RandomState(0						)
														SCREAMING_SNAKE_CASE__       :      Optional[Any]     =						pipe(
														    prompt=_UpperCAmelCase,			image=_UpperCAmelCase,			strength=0.75,			guidance_scale=7.5,			num_inference_steps=2_0,			generator=_UpperCAmelCase,			output_type="np",			)
														SCREAMING_SNAKE_CASE__       :      Optional[Any]     =						output.images
														SCREAMING_SNAKE_CASE__       :      List[str]     =						images[0, 2_5_5:2_5_8, 3_8_3:3_8_6, -1]

														assert images.shape == (1, 5_1_2, 7_6_8, 3)
														SCREAMING_SNAKE_CASE__       :      List[Any]     =						np.array([0.8043, 0.926, 0.9581, 0.8119, 0.8954, 0.913, 0.7209, 0.7463, 0.7431]						)
														# TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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













def _a (  SCREAMING_SNAKE_CASE__      :  int = 50_00_00_00       )	->    int:

							'''simple docstring'''






							SCREAMING_SNAKE_CASE__       :      str     =						set()
							SCREAMING_SNAKE_CASE__       :      Dict     =						int((limit - 24) ** (1 / 2)       )

							SCREAMING_SNAKE_CASE__       :      Any     =						set(range(3			,				prime_square_limit + 1			,				2       )       )
							primes.add(2       )
							for p in range(3			,				prime_square_limit + 1			,				2       ):
														if p not in primes:
																					continue
														primes.difference_update(set(range(p * p			,				prime_square_limit + 1			,				SCREAMING_SNAKE_CASE__       )       )       )

							for primea in primes:
														SCREAMING_SNAKE_CASE__       :      Optional[int]     =						primea * primea
														for primea in primes:
																					SCREAMING_SNAKE_CASE__       :      Union[str, Any]     =						primea * primea * primea
																					if square + cube >= limit - 16:
																												break
																					for primea in primes:
																												SCREAMING_SNAKE_CASE__       :      List[str]     =						primea * primea * primea * primea
																												SCREAMING_SNAKE_CASE__       :      Optional[int]     =						square + cube + tetr
																												if total >= limit:
																																			break
																												ret.add(SCREAMING_SNAKE_CASE__       )

							return len(SCREAMING_SNAKE_CASE__       )


if __name__ == "__main__":
		print(f"{solution() = }")






