



'''simple docstring'''





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        lowerCAmelCase_(				SCREAMING_SNAKE_CASE_						, unittest.TestCase     ):
  '''simple docstring'''




  __lowercase							:      Union[str, Any]				=      '''ssube/stable-diffusion-x4-upscaler-onnx'''




  def 							UpperCAmelCase_	(    self  ,__UpperCAmelCase=0      )    ->		List[str]:
      lowerCAmelCase__ :  Dict							      =					floats_tensor((1, 3, 128, 128)  ,rng=random.Random(__UpperCAmelCase      )      )
      lowerCAmelCase__ :  Union[str, Any]							      =					torch.manual_seed(__UpperCAmelCase      )
      lowerCAmelCase__ :  Dict							      =					{
          """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 							UpperCAmelCase_	(    self      )    ->		Optional[int]:
      lowerCAmelCase__ :  List[Any]							      =					OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint  ,provider="""CPUExecutionProvider"""      )
      pipe.set_progress_bar_config(disable=__UpperCAmelCase      )

      lowerCAmelCase__ :  List[Any]							      =					self.get_dummy_inputs()
      lowerCAmelCase__ :  Optional[Any]							      =					pipe(**__UpperCAmelCase      ).images
      lowerCAmelCase__ :  List[str]							      =					image[0, -3:, -3:, -1].flatten()

      # started as 128, should now be 512
      assert image.shape == (1, 512, 512, 3)
      lowerCAmelCase__ :  int							      =					np.array(
          [0.6_9_7_4_7_8_2, 0.6_8_9_0_2_0_9_3, 0.7_0_1_3_5_8_8_5, 0.7_5_8_3_6_1_8, 0.7_8_0_4_5_4_5, 0.7_8_5_4_9_1_2, 0.7_8_6_6_7_4_2_6, 0.7_8_7_4_3_8_6_3, 0.7_8_0_7_0_2_2_3]      )
      assert np.abs(image_slice - expected_slice      ).max() < 1E-1




  def 							UpperCAmelCase_	(    self      )    ->		int:
      lowerCAmelCase__ :  Optional[int]							      =					OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint  ,provider="""CPUExecutionProvider"""      )
      lowerCAmelCase__ :  Union[str, Any]							      =					PNDMScheduler.from_config(pipe.scheduler.config  ,skip_prk_steps=__UpperCAmelCase      )
      pipe.set_progress_bar_config(disable=__UpperCAmelCase      )

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

      assert image.shape == (1, 512, 512, 3)
      lowerCAmelCase__ :  Tuple							      =					np.array(
          [0.6_8_9_8_8_9_2, 0.5_9_2_4_0_5_5_6, 0.5_2_4_9_9_5_2_7, 0.5_8_8_6_6_2_1_5, 0.5_2_2_5_8_2_3_5, 0.5_2_5_7_2_7_1_5, 0.6_2_4_1_4_4_7_3, 0.6_1_7_4_3_8_7, 0.6_2_1_4_9_6_4]      )
      assert np.abs(image_slice.flatten() - expected_slice      ).max() < 1E-1




  def 							UpperCAmelCase_	(    self      )    ->		Optional[Any]:
      lowerCAmelCase__ :  List[Any]							      =					OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint  ,provider="""CPUExecutionProvider"""      )
      lowerCAmelCase__ :  str							      =					DPMSolverMultistepScheduler.from_config(pipe.scheduler.config      )
      pipe.set_progress_bar_config(disable=__UpperCAmelCase      )

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

      assert image.shape == (1, 512, 512, 3)
      lowerCAmelCase__ :  List[Any]							      =					np.array(
          [0.7_6_5_9_2_7_8, 0.7_6_4_3_7_6_6_4, 0.7_5_5_7_9_1_0_7, 0.7_6_9_1_1_1_6, 0.7_7_6_6_6_9_8_6, 0.7_7_2_7_6_7_2, 0.7_7_5_8_6_6_4, 0.7_8_1_2_2_2_6, 0.7_6_9_4_2_5_1_5]      )

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




  def 							UpperCAmelCase_	(    self      )    ->		int:
      lowerCAmelCase__ :  Optional[Any]							      =					OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint  ,provider="""CPUExecutionProvider"""      )
      lowerCAmelCase__ :  Optional[Any]							      =					EulerDiscreteScheduler.from_config(pipe.scheduler.config      )
      pipe.set_progress_bar_config(disable=__UpperCAmelCase      )

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

      assert image.shape == (1, 512, 512, 3)
      lowerCAmelCase__ :  int							      =					np.array(
          [0.6_9_7_4_7_8_2, 0.6_8_9_0_2_0_9_3, 0.7_0_1_3_5_8_8_5, 0.7_5_8_3_6_1_8, 0.7_8_0_4_5_4_5, 0.7_8_5_4_9_1_2, 0.7_8_6_6_7_4_2_6, 0.7_8_7_4_3_8_6_3, 0.7_8_0_7_0_2_2_3]      )
      assert np.abs(image_slice.flatten() - expected_slice      ).max() < 1E-1






  def 							UpperCAmelCase_	(    self      )    ->		str:
      lowerCAmelCase__ :  List[Any]							      =					OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint  ,provider="""CPUExecutionProvider"""      )
      lowerCAmelCase__ :  Optional[int]							      =					EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config      )
      pipe.set_progress_bar_config(disable=__UpperCAmelCase      )

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

      assert image.shape == (1, 512, 512, 3)
      lowerCAmelCase__ :  int							      =					np.array(
          [0.7_7_4_2_4_4_9_6, 0.7_7_3_6_0_1, 0.7_6_4_5_2_8_8, 0.7_7_6_9_5_9_8, 0.7_7_7_2_7_3_9, 0.7_7_3_8_6_8_8, 0.7_8_1_8_7_2_3_3, 0.7_7_8_7_9_5_8_4, 0.7_6_7_0_4_3]      )

      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 							UpperCAmelCase_	(    self      )    ->		Optional[Any]:
      return (
          "CUDAExecutionProvider",
          {
              "gpu_mem_limit": "15000000000",  # 15GB
              "arena_extend_strategy": "kSameAsRequested",
          },
      )




  @property
  def 							UpperCAmelCase_	(    self      )    ->		Any:
      lowerCAmelCase__ :  List[Any]							      =					ort.SessionOptions()
      lowerCAmelCase__ :  List[Any]							      =					False
      return options




  def 							UpperCAmelCase_	(    self      )    ->		List[str]:
      lowerCAmelCase__ :  Dict							      =					load_image(
          """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
          """/img2img/sketch-mountains-input.jpg"""      )
      lowerCAmelCase__ :  int							      =					init_image.resize((128, 128)      )
      # using the PNDM scheduler by default
      lowerCAmelCase__ :  Dict							      =					OnnxStableDiffusionUpscalePipeline.from_pretrained(
          """ssube/stable-diffusion-x4-upscaler-onnx"""  ,provider=self.gpu_provider  ,sess_options=self.gpu_options  ,)
      pipe.set_progress_bar_config(disable=__UpperCAmelCase      )

      lowerCAmelCase__ :  str							      =					"""A fantasy landscape, trending on artstation"""

      lowerCAmelCase__ :  Union[str, Any]							      =					torch.manual_seed(0      )
      lowerCAmelCase__ :  str							      =					pipe(
          prompt=__UpperCAmelCase  ,image=__UpperCAmelCase  ,guidance_scale=7.5  ,num_inference_steps=10  ,generator=__UpperCAmelCase  ,output_type="""np"""  ,)
      lowerCAmelCase__ :  Any							      =					output.images
      lowerCAmelCase__ :  Optional[int]							      =					images[0, 255:258, 383:386, -1]

      assert images.shape == (1, 512, 512, 3)
      lowerCAmelCase__ :  Optional[Any]							      =					np.array([0.4_8_8_3, 0.4_9_4_7, 0.4_9_8_0, 0.4_9_7_5, 0.4_9_8_2, 0.4_9_8_0, 0.5_0_0_0, 0.5_0_0_6, 0.4_9_7_2]      )
      # TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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






  def 							UpperCAmelCase_	(    self      )    ->		Optional[int]:
      lowerCAmelCase__ :  str							      =					load_image(
          """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
          """/img2img/sketch-mountains-input.jpg"""      )
      lowerCAmelCase__ :  Optional[int]							      =					init_image.resize((128, 128)      )
      lowerCAmelCase__ :  Tuple							      =					LMSDiscreteScheduler.from_pretrained(
          """ssube/stable-diffusion-x4-upscaler-onnx"""  ,subfolder="""scheduler"""      )
      lowerCAmelCase__ :  List[Any]							      =					OnnxStableDiffusionUpscalePipeline.from_pretrained(
          """ssube/stable-diffusion-x4-upscaler-onnx"""  ,scheduler=__UpperCAmelCase  ,provider=self.gpu_provider  ,sess_options=self.gpu_options  ,)
      pipe.set_progress_bar_config(disable=__UpperCAmelCase      )

      lowerCAmelCase__ :  str							      =					"""A fantasy landscape, trending on artstation"""

      lowerCAmelCase__ :  str							      =					torch.manual_seed(0      )
      lowerCAmelCase__ :  Dict							      =					pipe(
          prompt=__UpperCAmelCase  ,image=__UpperCAmelCase  ,guidance_scale=7.5  ,num_inference_steps=20  ,generator=__UpperCAmelCase  ,output_type="""np"""  ,)
      lowerCAmelCase__ :  Optional[int]							      =					output.images
      lowerCAmelCase__ :  List[Any]							      =					images[0, 255:258, 383:386, -1]

      assert images.shape == (1, 512, 512, 3)
      lowerCAmelCase__ :  Union[str, Any]							      =					np.array(
          [0.5_0_1_7_3_7_5_3, 0.5_0_2_2_3_3_5_6, 0.5_0_2_0_3_9, 0.5_0_2_3_3_0_3_6, 0.5_0_2_3_7_2_5, 0.5_0_2_2_6_0_1, 0.5_0_1_8_7_5_8, 0.5_0_2_3_4_0_8_5, 0.5_0_2_4_1_5_6_6]      )
      # TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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






"""simple docstring"""
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 (							__snake_case,				unittest.TestCase				):






		'''simple docstring'''

		lowerCamelCase__				=ShapEPipeline
		lowerCamelCase__				=['prompt']
		lowerCamelCase__				=['prompt']
		lowerCamelCase__				=[
		    'num_images_per_prompt',
		    'num_inference_steps',
		    'generator',
		    'latents',
		    'guidance_scale',
		    'frame_size',
		    'output_type',
		    'return_dict',
		]
		lowerCamelCase__				=False



		@property
		def 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

			return 32



		@property
		def 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

			return 32



		@property
		def 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

			return self.time_input_dim * 4



		@property
		def 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

			return 8



		@property
		def 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

			__snake_case	:						Optional[Any]						=				CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''					)
			return tokenizer



		@property
		def 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

			torch.manual_seed(0					)
			__snake_case	:						Optional[int]						=				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=10_00				,       )
			return CLIPTextModelWithProjection(a_					)



		@property
		def 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

			torch.manual_seed(0					)

			__snake_case	:						Optional[int]						=				{
			    '''num_attention_heads''': 2,
			    '''attention_head_dim''': 16,
			    '''embedding_dim''': self.time_input_dim,
			    '''num_embeddings''': 32,
			    '''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	:						Optional[int]						=				PriorTransformer(**a_					)
			return model



		@property
		def 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

			torch.manual_seed(0					)

			__snake_case	:						Optional[int]						=				{
			    '''param_shapes''': (
			        (self.renderer_dim, 93),
			        (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''': 12,
			    '''background''': (
			        0.1,
			        0.1,
			        0.1,
			    ),
			}
			__snake_case	:						List[Any]						=				ShapERenderer(**a_					)
			return model



		def 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

			__snake_case	:						int						=				self.dummy_prior
			__snake_case	:						str						=				self.dummy_text_encoder
			__snake_case	:						str						=				self.dummy_tokenizer
			__snake_case	:						Tuple						=				self.dummy_renderer

			__snake_case	:						int						=				HeunDiscreteScheduler(
			    beta_schedule='''exp'''				,       num_train_timesteps=10_24				,       prediction_type='''sample'''				,       use_karras_sigmas=a_				,       clip_sample=a_				,       clip_sample_range=1.0				,       )
			__snake_case	:						Union[str, Any]						=				{
			    '''prior''': prior,
			    '''text_encoder''': text_encoder,
			    '''tokenizer''': tokenizer,
			    '''renderer''': renderer,
			    '''scheduler''': scheduler,
			}

			return components



		def 				SCREAMING_SNAKE_CASE  (self				,       a_				,       a_=0					):





			'''simple docstring'''

			if str(a_					).startswith('''mps'''					):
				__snake_case	:						Tuple						=				torch.manual_seed(a_					)
			else:
				__snake_case	:						Union[str, Any]						=				torch.Generator(device=a_					).manual_seed(a_					)
			__snake_case	:						Optional[int]						=				{
			    '''prompt''': '''horse''',
			    '''generator''': generator,
			    '''num_inference_steps''': 1,
			    '''frame_size''': 32,
			    '''output_type''': '''np''',
			}
			return inputs



		def 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

			__snake_case	:						Optional[int]						=				'''cpu'''

			__snake_case	:						str						=				self.get_dummy_components()

			__snake_case	:						List[Any]						=				self.pipeline_class(**a_					)
			__snake_case	:						str						=				pipe.to(a_					)

			pipe.set_progress_bar_config(disable=a_					)

			__snake_case	:						Optional[Any]						=				pipe(**self.get_dummy_inputs(a_					)					)
			__snake_case	:						List[str]						=				output.images[0]
			__snake_case	:						List[str]						=				image[0, -3:, -3:, -1]

			assert image.shape == (20, 32, 32, 3)

			__snake_case	:						List[Any]						=				np.array(
			    [
			        0.0003_9216,
			        0.0003_9216,
			        0.0003_9216,
			        0.0003_9216,
			        0.0003_9216,
			        0.0003_9216,
			        0.0003_9216,
			        0.0003_9216,
			        0.0003_9216,
			    ]					)

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



		def 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

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



		def 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

			__snake_case	:						Any						=				torch_device == '''cpu'''
			__snake_case	:						Tuple						=				True

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



		def 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

			__snake_case	:						int						=				self.get_dummy_components()
			__snake_case	:						int						=				self.pipeline_class(**a_					)
			__snake_case	:						int						=				pipe.to(a_					)
			pipe.set_progress_bar_config(disable=a_					)

			__snake_case	:						Optional[Any]						=				1
			__snake_case	:						List[Any]						=				2

			__snake_case	:						int						=				self.get_dummy_inputs(a_					)

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

			__snake_case	:						Union[str, Any]						=				pipe(**a_				,       num_images_per_prompt=a_					)[0]

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



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






		'''simple docstring'''

		def 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

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



		def 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

			__snake_case	:						List[str]						=				load_numpy(
			    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
			    '''/shap_e/test_shap_e_np_out.npy'''					)
			__snake_case	:						Optional[Any]						=				ShapEPipeline.from_pretrained('''openai/shap-e'''					)
			__snake_case	:						Optional[int]						=				pipe.to(a_					)
			pipe.set_progress_bar_config(disable=a_					)

			__snake_case	:						Optional[Any]						=				torch.Generator(device=a_					).manual_seed(0					)

			__snake_case	:						str						=				pipe(
			    '''a shark'''				,       generator=a_				,       guidance_scale=15.0				,       num_inference_steps=64				,       frame_size=64				,       output_type='''np'''				,       ).images[0]

			assert images.shape == (20, 64, 64, 3)

			assert_mean_pixel_difference(a_				,       a_					)





