




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     lowerCamelCase__(					__lowerCamelCase ,     unittest.TestCase):
      UpperCAmelCase__      :	Dict    		=   ShapEPipeline
      UpperCAmelCase__      :	Any    		=   ['prompt']
      UpperCAmelCase__      :	List[Any]    		=   ['prompt']
      UpperCAmelCase__      :	str    		=   [
          'num_images_per_prompt',
          'num_inference_steps',
          'generator',
          'latents',
          'guidance_scale',
          'frame_size',
          'output_type',
          'return_dict',
      ]
      UpperCAmelCase__      :	Dict    		=   False
      @property
      def 			lowerCAmelCase__						(				self: int   ):
             return 32
      @property
      def 			lowerCAmelCase__						(				self: Union[str, Any]   ):
             return 32
      @property
      def 			lowerCAmelCase__						(				self: int   ):
             return self.time_input_dim * 4
      @property
      def 			lowerCAmelCase__						(				self: Optional[Any]   ):
             return 8
      @property
      def 			lowerCAmelCase__						(				self: List[str]   ):
             __lowerCamelCase   =						CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""   )
             return tokenizer
      @property
      def 			lowerCAmelCase__						(				self: str   ):
             torch.manual_seed(0   )
             __lowerCamelCase   =						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(UpperCamelCase_   )
      @property
      def 			lowerCAmelCase__						(				self: List[str]   ):
             torch.manual_seed(0   )

             __lowerCamelCase   =						{
                 """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,
             }

             __lowerCamelCase   =						PriorTransformer(**UpperCamelCase_   )
             return model
      @property
      def 			lowerCAmelCase__						(				self: Dict   ):
             torch.manual_seed(0   )

             __lowerCamelCase   =						{
                 """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,
                 ),
             }
             __lowerCamelCase   =						ShapERenderer(**UpperCamelCase_   )
             return model
      def 			lowerCAmelCase__						(				self: str   ):
             __lowerCamelCase   =						self.dummy_prior
             __lowerCamelCase   =						self.dummy_text_encoder
             __lowerCamelCase   =						self.dummy_tokenizer
             __lowerCamelCase   =						self.dummy_renderer

             __lowerCamelCase   =						HeunDiscreteScheduler(
                 beta_schedule="""exp"""    ,       num_train_timesteps=10_24    ,       prediction_type="""sample"""    ,       use_karras_sigmas=UpperCamelCase_    ,       clip_sample=UpperCamelCase_    ,       clip_sample_range=1.0    ,       )
             __lowerCamelCase   =						{
                 """prior""": prior,
                 """text_encoder""": text_encoder,
                 """tokenizer""": tokenizer,
                 """renderer""": renderer,
                 """scheduler""": scheduler,
             }

             return components
      def 			lowerCAmelCase__						(				self: str    ,       UpperCamelCase_: Dict    ,       UpperCamelCase_: Optional[Any]=0   ):
             if str(UpperCamelCase_   ).startswith("""mps"""   ):
                    __lowerCamelCase   =						torch.manual_seed(UpperCamelCase_   )
             else:
                    __lowerCamelCase   =						torch.Generator(device=UpperCamelCase_   ).manual_seed(UpperCamelCase_   )
             __lowerCamelCase   =						{
                 """prompt""": """horse""",
                 """generator""": generator,
                 """num_inference_steps""": 1,
                 """frame_size""": 32,
                 """output_type""": """np""",
             }
             return inputs
      def 			lowerCAmelCase__						(				self: List[str]   ):
             __lowerCamelCase   =						"""cpu"""

             __lowerCamelCase   =						self.get_dummy_components()

             __lowerCamelCase   =						self.pipeline_class(**UpperCamelCase_   )
             __lowerCamelCase   =						pipe.to(UpperCamelCase_   )

             pipe.set_progress_bar_config(disable=UpperCamelCase_   )

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

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

             __lowerCamelCase   =						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 			lowerCAmelCase__						(				self: List[str]   ):
             # NOTE: Larger batch sizes cause this test to timeout, only test on smaller batches
             self._test_inference_batch_consistent(batch_sizes=[1, 2]   )
      def 			lowerCAmelCase__						(				self: Optional[int]   ):
             __lowerCamelCase   =						torch_device == """cpu"""
             __lowerCamelCase   =						True

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




      def 			lowerCAmelCase__						(				self: Union[str, Any]   ):
             __lowerCamelCase   =						self.get_dummy_components()
             __lowerCamelCase   =						self.pipeline_class(**UpperCamelCase_   )
             __lowerCamelCase   =						pipe.to(UpperCamelCase_   )
             pipe.set_progress_bar_config(disable=UpperCamelCase_   )

             __lowerCamelCase   =						1
             __lowerCamelCase   =						2

             __lowerCamelCase   =						self.get_dummy_inputs(UpperCamelCase_   )

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

             __lowerCamelCase   =						pipe(**UpperCamelCase_    ,       num_images_per_prompt=UpperCamelCase_   )[0]

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





@slow
@require_torch_gpu
class     lowerCamelCase__(					unittest.TestCase):
      def 			lowerCAmelCase__						(				self: List[Any]   ):
             # clean up the VRAM after each test
             super().tearDown()
             gc.collect()
             torch.cuda.empty_cache()




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

             __lowerCamelCase   =						torch.Generator(device=UpperCamelCase_   ).manual_seed(0   )

             __lowerCamelCase   =						pipe(
                 """a shark"""    ,       generator=UpperCamelCase_    ,       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(UpperCamelCase_    ,       UpperCamelCase_   )








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 		SCREAMING_SNAKE_CASE__				(       _UpperCAmelCase						,      unittest.TestCase  ):
      A_						:				Tuple							=    'hf-internal-testing/tiny-random-OnnxStableDiffusionPipeline'


      def    a (self       :      int				,			a__       :      List[Any]=0      ):





         """simple docstring"""



         __snake_case							       =   floats_tensor((1, 3, 128, 128)				,			rng=random.Random(a__      )      )
         __snake_case							       =   np.random.RandomState(a__      )
         __snake_case							       =   {
             '''prompt''': '''A painting of a squirrel eating a burger''',
             '''image''': image,
             '''generator''': generator,
             '''num_inference_steps''': 3,
             '''strength''': 0.7_5,
             '''guidance_scale''': 7.5,
             '''output_type''': '''numpy''',
         }
         return inputs


      def    a (self       :      List[Any]      ):





         """simple docstring"""



         __snake_case							       =   OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint				,			provider='''CPUExecutionProvider'''      )
         pipe.set_progress_bar_config(disable=a__      )

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

         assert image.shape == (1, 128, 128, 3)
         __snake_case							       =   np.array([0.6_9_6_4_3, 0.5_8_4_8_4, 0.5_0_3_1_4, 0.5_8_7_6_0, 0.5_5_3_6_8, 0.5_9_6_4_3, 0.5_1_5_2_9, 0.4_1_2_1_7, 0.4_9_0_8_7]      )
         assert np.abs(image_slice - expected_slice      ).max() < 1E-1


      def    a (self       :      Dict      ):





         """simple docstring"""



         __snake_case							       =   OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint				,			provider='''CPUExecutionProvider'''      )
         __snake_case							       =   PNDMScheduler.from_config(pipe.scheduler.config				,			skip_prk_steps=a__      )
         pipe.set_progress_bar_config(disable=a__      )

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

         assert image.shape == (1, 128, 128, 3)
         __snake_case							       =   np.array([0.6_1_7_3_7, 0.5_4_6_4_2, 0.5_3_1_8_3, 0.5_4_4_6_5, 0.5_2_7_4_2, 0.6_0_5_2_5, 0.4_9_9_6_9, 0.4_0_6_5_5, 0.4_8_1_5_4]      )

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


      def    a (self       :      List[str]      ):





         """simple docstring"""



         __snake_case							       =   OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint				,			provider='''CPUExecutionProvider'''      )
         __snake_case							       =   LMSDiscreteScheduler.from_config(pipe.scheduler.config      )
         pipe.set_progress_bar_config(disable=a__      )

         # warmup pass to apply optimizations
         __snake_case							       =   pipe(**self.get_dummy_inputs()      )

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

         assert image.shape == (1, 128, 128, 3)
         __snake_case							       =   np.array([0.5_2_7_6_1, 0.5_9_9_7_7, 0.4_9_0_3_3, 0.4_9_6_1_9, 0.5_4_2_8_2, 0.5_0_3_1_1, 0.4_7_6_0_0, 0.4_0_9_1_8, 0.4_5_2_0_3]      )

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


      def    a (self       :      Any      ):





         """simple docstring"""



         __snake_case							       =   OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint				,			provider='''CPUExecutionProvider'''      )
         __snake_case							       =   EulerDiscreteScheduler.from_config(pipe.scheduler.config      )
         pipe.set_progress_bar_config(disable=a__      )

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

         assert image.shape == (1, 128, 128, 3)
         __snake_case							       =   np.array([0.5_2_9_1_1, 0.6_0_0_0_4, 0.4_9_2_2_9, 0.4_9_8_0_5, 0.5_4_5_0_2, 0.5_0_6_8_0, 0.4_7_7_7_7, 0.4_1_0_2_8, 0.4_5_3_0_4]      )

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


      def    a (self       :      Dict      ):





         """simple docstring"""



         __snake_case							       =   OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint				,			provider='''CPUExecutionProvider'''      )
         __snake_case							       =   EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config      )
         pipe.set_progress_bar_config(disable=a__      )

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

         assert image.shape == (1, 128, 128, 3)
         __snake_case							       =   np.array([0.5_2_9_1_1, 0.6_0_0_0_4, 0.4_9_2_2_9, 0.4_9_8_0_5, 0.5_4_5_0_2, 0.5_0_6_8_0, 0.4_7_7_7_7, 0.4_1_0_2_8, 0.4_5_3_0_4]      )

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


      def    a (self       :      List[str]      ):





         """simple docstring"""



         __snake_case							       =   OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint				,			provider='''CPUExecutionProvider'''      )
         __snake_case							       =   DPMSolverMultistepScheduler.from_config(pipe.scheduler.config      )
         pipe.set_progress_bar_config(disable=a__      )

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

         assert image.shape == (1, 128, 128, 3)
         __snake_case							       =   np.array([0.6_5_3_3_1, 0.5_8_2_7_7, 0.4_8_2_0_4, 0.5_6_0_5_9, 0.5_3_6_6_5, 0.5_6_2_3_5, 0.5_0_9_6_9, 0.4_0_0_0_9, 0.4_6_5_5_2]      )

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






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


      @property
      def    a (self       :      List[str]      ):





         """simple docstring"""



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


      @property
      def    a (self       :      Optional[Any]      ):





         """simple docstring"""



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


      def    a (self       :      Optional[Any]      ):





         """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((768, 512)      )
         # using the PNDM scheduler by default
         __snake_case							       =   OnnxStableDiffusionImgaImgPipeline.from_pretrained(
             '''CompVis/stable-diffusion-v1-4'''				,			revision='''onnx'''				,			safety_checker=a__				,			feature_extractor=a__				,			provider=self.gpu_provider				,			sess_options=self.gpu_options				,			)
         pipe.set_progress_bar_config(disable=a__      )

         __snake_case							       =   '''A fantasy landscape, trending on artstation'''

         __snake_case							       =   np.random.RandomState(0      )
         __snake_case							       =   pipe(
             prompt=a__				,			image=a__				,			strength=0.7_5				,			guidance_scale=7.5				,			num_inference_steps=10				,			generator=a__				,			output_type='''np'''				,			)
         __snake_case							       =   output.images
         __snake_case							       =   images[0, 255:258, 383:386, -1]

         assert images.shape == (1, 512, 768, 3)
         __snake_case							       =   np.array([0.4_9_0_9, 0.5_0_5_9, 0.5_3_7_2, 0.4_6_2_3, 0.4_8_7_6, 0.5_0_4_9, 0.4_8_2_0, 0.4_9_5_6, 0.5_0_1_9]      )
         # 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       :      Dict      ):





         """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((768, 512)      )
         __snake_case							       =   LMSDiscreteScheduler.from_pretrained(
             '''runwayml/stable-diffusion-v1-5'''				,			subfolder='''scheduler'''				,			revision='''onnx'''      )
         __snake_case							       =   OnnxStableDiffusionImgaImgPipeline.from_pretrained(
             '''runwayml/stable-diffusion-v1-5'''				,			revision='''onnx'''				,			scheduler=a__				,			safety_checker=a__				,			feature_extractor=a__				,			provider=self.gpu_provider				,			sess_options=self.gpu_options				,			)
         pipe.set_progress_bar_config(disable=a__      )

         __snake_case							       =   '''A fantasy landscape, trending on artstation'''

         __snake_case							       =   np.random.RandomState(0      )
         __snake_case							       =   pipe(
             prompt=a__				,			image=a__				,			strength=0.7_5				,			guidance_scale=7.5				,			num_inference_steps=20				,			generator=a__				,			output_type='''np'''				,			)
         __snake_case							       =   output.images
         __snake_case							       =   images[0, 255:258, 383:386, -1]

         assert images.shape == (1, 512, 768, 3)
         __snake_case							       =   np.array([0.8_0_4_3, 0.9_2_6, 0.9_5_8_1, 0.8_1_1_9, 0.8_9_5_4, 0.9_1_3, 0.7_2_0_9, 0.7_4_6_3, 0.7_4_3_1]      )
         # TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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

