


import gc
import random
import unittest

import numpy as np
import torch
from transformers import CLIPImageProcessor, CLIPVisionConfig, CLIPVisionModel

from diffusers import HeunDiscreteScheduler, PriorTransformer, ShapEImgaImgPipeline
from diffusers.pipelines.shap_e import ShapERenderer
from diffusers.utils import floats_tensor, load_image, 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			(_a				,      unittest.TestCase						):
       _lowercase											=     ShapEImgaImgPipeline
       _lowercase											=     ["""image"""]
       _lowercase											=     ["""image"""]
       _lowercase											=     [
           """num_images_per_prompt""",
           """num_inference_steps""",
           """generator""",
           """latents""",
           """guidance_scale""",
           """frame_size""",
           """output_type""",
           """return_dict""",
       ]
       _lowercase											=     False







       @property
       def 		snake_case_	(			self: Dict					):






            '''simple docstring'''




            return 32







       @property
       def 		snake_case_	(			self: Any					):






            '''simple docstring'''




            return 32







       @property
       def 		snake_case_	(			self: Any					):






            '''simple docstring'''




            return self.time_input_dim * 4







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






            '''simple docstring'''




            return 8







       @property
       def 		snake_case_	(			self: List[Any]					):






            '''simple docstring'''




            torch.manual_seed(0					)
            __UpperCamelCase =			CLIPVisionConfig(
                hidden_size=self.text_embedder_hidden_size,image_size=64,projection_dim=self.text_embedder_hidden_size,intermediate_size=37,num_attention_heads=4,num_channels=3,num_hidden_layers=5,patch_size=1,)

            __UpperCamelCase =			CLIPVisionModel(A_					)
            return model







       @property
       def 		snake_case_	(			self: Tuple					):






            '''simple docstring'''




            __UpperCamelCase =			CLIPImageProcessor(
                crop_size=224,do_center_crop=A_,do_normalize=A_,do_resize=A_,image_mean=[0.4_8_1_4_5_4_6_6, 0.4_5_7_8_2_7_5, 0.4_0_8_2_1_0_7_3],image_std=[0.2_6_8_6_2_9_5_4, 0.2_6_1_3_0_2_5_8, 0.2_7_5_7_7_7_1_1],resample=3,size=224,)

            return image_processor







       @property
       def 		snake_case_	(			self: int					):






            '''simple docstring'''




            torch.manual_seed(0					)

            __UpperCamelCase =			{
                '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',
                'embedding_proj_norm_type': 'layer',
                'encoder_hid_proj_type': None,
                'added_emb_type': None,
            }

            __UpperCamelCase =			PriorTransformer(**A_					)
            return model







       @property
       def 		snake_case_	(			self: List[Any]					):






            '''simple docstring'''




            torch.manual_seed(0					)

            __UpperCamelCase =			{
                '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,
                ),
            }
            __UpperCamelCase =			ShapERenderer(**A_					)
            return model







       def 		snake_case_	(			self: Optional[Any]					):






            '''simple docstring'''




            __UpperCamelCase =			self.dummy_prior
            __UpperCamelCase =			self.dummy_image_encoder
            __UpperCamelCase =			self.dummy_image_processor
            __UpperCamelCase =			self.dummy_renderer

            __UpperCamelCase =			HeunDiscreteScheduler(
                beta_schedule='exp',num_train_timesteps=1024,prediction_type='sample',use_karras_sigmas=A_,clip_sample=A_,clip_sample_range=1.0,)
            __UpperCamelCase =			{
                'prior': prior,
                'image_encoder': image_encoder,
                'image_processor': image_processor,
                'renderer': renderer,
                'scheduler': scheduler,
            }

            return components







       def 		snake_case_	(			self: Union[str, Any],A_: Optional[Any],A_: str=0					):






            '''simple docstring'''




            __UpperCamelCase =			floats_tensor((1, 3, 64, 64),rng=random.Random(A_					)					).to(A_					)

            if str(A_					).startswith('mps'					):
                 __UpperCamelCase =			torch.manual_seed(A_					)
            else:
                 __UpperCamelCase =			torch.Generator(device=A_					).manual_seed(A_					)
            __UpperCamelCase =			{
                'image': input_image,
                'generator': generator,
                'num_inference_steps': 1,
                'frame_size': 32,
                'output_type': 'np',
            }
            return inputs







       def 		snake_case_	(			self: List[Any]					):






            '''simple docstring'''




            __UpperCamelCase =			'cpu'

            __UpperCamelCase =			self.get_dummy_components()

            __UpperCamelCase =			self.pipeline_class(**A_					)
            __UpperCamelCase =			pipe.to(A_					)

            pipe.set_progress_bar_config(disable=A_					)

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

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

            __UpperCamelCase =			np.array(
                [
                    0.0_0_0_3_9_2_1_6,
                    0.0_0_0_3_9_2_1_6,
                    0.0_0_0_3_9_2_1_6,
                    0.0_0_0_3_9_2_1_6,
                    0.0_0_0_3_9_2_1_6,
                    0.0_0_0_3_9_2_1_6,
                    0.0_0_0_3_9_2_1_6,
                    0.0_0_0_3_9_2_1_6,
                    0.0_0_0_3_9_2_1_6,
                ]					)

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







       def 		snake_case_	(			self: Any					):






            '''simple docstring'''




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







       def 		snake_case_	(			self: Optional[Any]					):






            '''simple docstring'''




            __UpperCamelCase =			torch_device == 'cpu'
            __UpperCamelCase =			True
            self._test_inference_batch_single_identical(
                batch_size=2,test_max_difference=A_,relax_max_difference=A_,)







       def 		snake_case_	(			self: Union[str, Any]					):






            '''simple docstring'''




            __UpperCamelCase =			self.get_dummy_components()
            __UpperCamelCase =			self.pipeline_class(**A_					)
            __UpperCamelCase =			pipe.to(A_					)
            pipe.set_progress_bar_config(disable=A_					)

            __UpperCamelCase =			1
            __UpperCamelCase =			2

            __UpperCamelCase =			self.get_dummy_inputs(A_					)

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

            __UpperCamelCase =			pipe(**A_,num_images_per_prompt=A_					)[0]

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



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







       def 		snake_case_	(			self: List[Any]					):






            '''simple docstring'''




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







       def 		snake_case_	(			self: Any					):






            '''simple docstring'''




            __UpperCamelCase =			load_image(
                'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main' '/shap_e/corgi.png'					)
            __UpperCamelCase =			load_numpy(
                'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
                '/shap_e/test_shap_e_img2img_out.npy'					)
            __UpperCamelCase =			ShapEImgaImgPipeline.from_pretrained('openai/shap-e-img2img'					)
            __UpperCamelCase =			pipe.to(A_					)
            pipe.set_progress_bar_config(disable=A_					)

            __UpperCamelCase =			torch.Generator(device=A_					).manual_seed(0					)

            __UpperCamelCase =			pipe(
                A_,generator=A_,guidance_scale=3.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_					)










"""simple docstring"""


import unittest

import numpy as np

from diffusers import OnnxStableDiffusionInpaintPipelineLegacy
from diffusers.utils.testing_utils import (
    is_onnx_available,
    load_image,
    load_numpy,
    nightly,
    require_onnxruntime,
    require_torch_gpu,
)


if is_onnx_available():
		import onnxruntime as ort




@nightly
@require_onnxruntime
@require_torch_gpu
class 						__UpperCamelCase     (     unittest.TestCase						):
					@property
					def        UpperCAmelCase__ (							self							:	Union[str, Any]   )	->       Dict:
											return (
											    "CUDAExecutionProvider",
											    {
											        "gpu_mem_limit": "15000000000",  # 15GB
											        "arena_extend_strategy": "kSameAsRequested",
											    },
											)
					@property
					def        UpperCAmelCase__ (							self							:	int   )	->       Optional[Any]:
											lowerCAmelCase   :Any						=			ort.SessionOptions()
											lowerCAmelCase   :Dict						=			False
											return options
					def        UpperCAmelCase__ (							self							:	Tuple   )	->       str:
											lowerCAmelCase   :Union[str, Any]						=			load_image(
											    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
											    '/in_paint/overture-creations-5sI6fQgYIuo.png'   )
											lowerCAmelCase   :str						=			load_image(
											    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
											    '/in_paint/overture-creations-5sI6fQgYIuo_mask.png'   )
											lowerCAmelCase   :List[Any]						=			load_numpy(
											    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
											    '/in_paint/red_cat_sitting_on_a_park_bench_onnx.npy'   )

											# using the PNDM scheduler by default
											lowerCAmelCase   :Tuple						=			OnnxStableDiffusionInpaintPipelineLegacy.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   )

											lowerCAmelCase   :Optional[Any]						=			'A red cat sitting on a park bench'

											lowerCAmelCase   :Optional[Any]						=			np.random.RandomState(0   )
											lowerCAmelCase   :Union[str, Any]						=			pipe(
											    prompt=UpperCAmelCase     , image=UpperCAmelCase     , mask_image=UpperCAmelCase     , strength=0.7_5     , guidance_scale=7.5     , num_inference_steps=15     , generator=UpperCAmelCase     , output_type='np'     , )

											lowerCAmelCase   :Optional[int]						=			output.images[0]

											assert image.shape == (512, 512, 3)
											assert np.abs(expected_image - image   ).max() < 1e-2