




'''simple docstring'''



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 				_a							(       _UpperCAmelCase      ,	unittest.TestCase      ):





       '''simple docstring'''






       A		:							Tuple								=		ShapEImgaImgPipeline
       A		:							List[str]								=		['''image''']
       A		:							Tuple								=		['''image''']
       A		:							Optional[int]								=		[
           '''num_images_per_prompt''',
           '''num_inference_steps''',
           '''generator''',
           '''latents''',
           '''guidance_scale''',
           '''frame_size''',
           '''output_type''',
           '''return_dict''',
       ]
       A		:							Dict								=		False







       @property
       def     UpperCamelCase_				(	self   ):







             '''simple docstring'''







             return 32







       @property
       def     UpperCamelCase_				(	self   ):







             '''simple docstring'''







             return 32







       @property
       def     UpperCamelCase_				(	self   ):







             '''simple docstring'''







             return self.time_input_dim * 4







       @property
       def     UpperCamelCase_				(	self   ):







             '''simple docstring'''







             return 8







       @property
       def     UpperCamelCase_				(	self   ):







             '''simple docstring'''







             torch.manual_seed(0   )
             SCREAMING_SNAKE_CASE					:	str   =		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,				)

             SCREAMING_SNAKE_CASE					:	List[Any]   =		CLIPVisionModel(A   )
             return model







       @property
       def     UpperCamelCase_				(	self   ):







             '''simple docstring'''







             SCREAMING_SNAKE_CASE					:	Optional[int]   =		CLIPImageProcessor(
                 crop_size=224,				do_center_crop=A,				do_normalize=A,				do_resize=A,				image_mean=[0.48_14_54_66, 0.4_57_82_75, 0.40_82_10_73],				image_std=[0.26_86_29_54, 0.26_13_02_58, 0.27_57_77_11],				resample=3,				size=224,				)

             return image_processor







       @property
       def     UpperCamelCase_				(	self   ):







             '''simple docstring'''







             torch.manual_seed(0   )

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

             SCREAMING_SNAKE_CASE					:	str   =		PriorTransformer(**A   )
             return model







       @property
       def     UpperCamelCase_				(	self   ):







             '''simple docstring'''







             torch.manual_seed(0   )

             SCREAMING_SNAKE_CASE					:	Union[str, Any]   =		{
                 '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,
                 ),
             }
             SCREAMING_SNAKE_CASE					:	Tuple   =		ShapERenderer(**A   )
             return model







       def     UpperCamelCase_				(	self   ):







             '''simple docstring'''







             SCREAMING_SNAKE_CASE					:	Tuple   =		self.dummy_prior
             SCREAMING_SNAKE_CASE					:	List[Any]   =		self.dummy_image_encoder
             SCREAMING_SNAKE_CASE					:	Optional[int]   =		self.dummy_image_processor
             SCREAMING_SNAKE_CASE					:	Any   =		self.dummy_renderer

             SCREAMING_SNAKE_CASE					:	List[Any]   =		HeunDiscreteScheduler(
                 beta_schedule='exp',				num_train_timesteps=1_024,				prediction_type='sample',				use_karras_sigmas=A,				clip_sample=A,				clip_sample_range=1.0,				)
             SCREAMING_SNAKE_CASE					:	str   =		{
                 'prior': prior,
                 'image_encoder': image_encoder,
                 'image_processor': image_processor,
                 'renderer': renderer,
                 'scheduler': scheduler,
             }

             return components







       def     UpperCamelCase_				(	self,				A,				A=0   ):







             '''simple docstring'''







             SCREAMING_SNAKE_CASE					:	int   =		floats_tensor((1, 3, 64, 64),				rng=random.Random(A   )   ).to(A   )

             if str(A   ).startswith('mps'   ):
                   SCREAMING_SNAKE_CASE					:	Optional[int]   =		torch.manual_seed(A   )
             else:
                   SCREAMING_SNAKE_CASE					:	Optional[int]   =		torch.Generator(device=A   ).manual_seed(A   )
             SCREAMING_SNAKE_CASE					:	int   =		{
                 'image': input_image,
                 'generator': generator,
                 'num_inference_steps': 1,
                 'frame_size': 32,
                 'output_type': 'np',
             }
             return inputs







       def     UpperCamelCase_				(	self   ):







             '''simple docstring'''







             SCREAMING_SNAKE_CASE					:	Union[str, Any]   =		'cpu'

             SCREAMING_SNAKE_CASE					:	List[Any]   =		self.get_dummy_components()

             SCREAMING_SNAKE_CASE					:	List[Any]   =		self.pipeline_class(**A   )
             SCREAMING_SNAKE_CASE					:	Union[str, Any]   =		pipe.to(A   )

             pipe.set_progress_bar_config(disable=A   )

             SCREAMING_SNAKE_CASE					:	int   =		pipe(**self.get_dummy_inputs(A   )   )
             SCREAMING_SNAKE_CASE					:	List[str]   =		output.images[0]
             SCREAMING_SNAKE_CASE					:	Any   =		image[0, -3:, -3:, -1]

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

             SCREAMING_SNAKE_CASE					:	Any   =		np.array(
                 [
                     0.00_03_92_16,
                     0.00_03_92_16,
                     0.00_03_92_16,
                     0.00_03_92_16,
                     0.00_03_92_16,
                     0.00_03_92_16,
                     0.00_03_92_16,
                     0.00_03_92_16,
                     0.00_03_92_16,
                 ]   )

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







       def     UpperCamelCase_				(	self   ):







             '''simple docstring'''







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







       def     UpperCamelCase_				(	self   ):







             '''simple docstring'''







             SCREAMING_SNAKE_CASE					:	Any   =		torch_device == 'cpu'
             SCREAMING_SNAKE_CASE					:	Union[str, Any]   =		True
             self._test_inference_batch_single_identical(
                 batch_size=2,				test_max_difference=A,				relax_max_difference=A,				)







       def     UpperCamelCase_				(	self   ):







             '''simple docstring'''







             SCREAMING_SNAKE_CASE					:	List[Any]   =		self.get_dummy_components()
             SCREAMING_SNAKE_CASE					:	str   =		self.pipeline_class(**A   )
             SCREAMING_SNAKE_CASE					:	Any   =		pipe.to(A   )
             pipe.set_progress_bar_config(disable=A   )

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

             SCREAMING_SNAKE_CASE					:	str   =		self.get_dummy_inputs(A   )

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

             SCREAMING_SNAKE_CASE					:	Tuple   =		pipe(**A,				num_images_per_prompt=A   )[0]

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






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





       '''simple docstring'''






       def     UpperCamelCase_				(	self   ):







             '''simple docstring'''







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







       def     UpperCamelCase_				(	self   ):







             '''simple docstring'''







             SCREAMING_SNAKE_CASE					:	int   =		load_image(
                 'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main' '/shap_e/corgi.png'   )
             SCREAMING_SNAKE_CASE					:	int   =		load_numpy(
                 'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
                 '/shap_e/test_shap_e_img2img_out.npy'   )
             SCREAMING_SNAKE_CASE					:	List[str]   =		ShapEImgaImgPipeline.from_pretrained('openai/shap-e-img2img'   )
             SCREAMING_SNAKE_CASE					:	Optional[Any]   =		pipe.to(A   )
             pipe.set_progress_bar_config(disable=A   )

             SCREAMING_SNAKE_CASE					:	Dict   =		torch.Generator(device=A   ).manual_seed(0   )

             SCREAMING_SNAKE_CASE					:	Union[str, Any]   =		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"""





# This is the module that test_patching.py uses to test patch_submodule()

import os  # noqa: this is just for tests
import os as renamed_os  # noqa: this is just for tests
from os import path  # noqa: this is just for tests
from os import path as renamed_path  # noqa: this is just for tests
from os.path import join  # noqa: this is just for tests
from os.path import join as renamed_join  # noqa: this is just for tests


SCREAMING_SNAKE_CASE__					=   open  # noqa: we just need to have a builtin inside this module to test it properly






