
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( a  ,   unittest.TestCase			):
				snake_case_						   =    ShapEImgaImgPipeline
				snake_case_						   =    ['''image''']
				snake_case_						   =    ['''image''']
				snake_case_						   =    [
				    '''num_images_per_prompt''',
				    '''num_inference_steps''',
				    '''generator''',
				    '''latents''',
				    '''guidance_scale''',
				    '''frame_size''',
				    '''output_type''',
				    '''return_dict''',
				]
				snake_case_						   =    False




				@property
				def 				SCREAMING_SNAKE_CASE_		(     self				)    ->      List[Any]:


								'''simple docstring'''




								return 32




				@property
				def 				SCREAMING_SNAKE_CASE_		(     self				)    ->      Dict:


								'''simple docstring'''




								return 32




				@property
				def 				SCREAMING_SNAKE_CASE_		(     self				)    ->      Optional[int]:


								'''simple docstring'''




								return self.time_input_dim * 4




				@property
				def 				SCREAMING_SNAKE_CASE_		(     self				)    ->      List[str]:


								'''simple docstring'''




								return 8




				@property
				def 				SCREAMING_SNAKE_CASE_		(     self				)    ->      Optional[int]:


								'''simple docstring'''




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

								__a										=    CLIPVisionModel(_snake_case				)
								return model




				@property
				def 				SCREAMING_SNAKE_CASE_		(     self				)    ->      List[Any]:


								'''simple docstring'''




								__a										=    CLIPImageProcessor(
								    crop_size=224		,     do_center_crop=_snake_case		,     do_normalize=_snake_case		,     do_resize=_snake_case		,     image_mean=[0.4814_5466, 0.457_8275, 0.4082_1073]		,     image_std=[0.2686_2954, 0.2613_0258, 0.2757_7711]		,     resample=3		,     size=224		,     )

								return image_processor




				@property
				def 				SCREAMING_SNAKE_CASE_		(     self				)    ->      Tuple:


								'''simple docstring'''




								torch.manual_seed(0				)

								__a										=    {
								    '''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,
								}

								__a										=    PriorTransformer(**_snake_case				)
								return model




				@property
				def 				SCREAMING_SNAKE_CASE_		(     self				)    ->      Tuple:


								'''simple docstring'''




								torch.manual_seed(0				)

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




				def 				SCREAMING_SNAKE_CASE_		(     self				)    ->      str:


								'''simple docstring'''




								__a										=    self.dummy_prior
								__a										=    self.dummy_image_encoder
								__a										=    self.dummy_image_processor
								__a										=    self.dummy_renderer

								__a										=    HeunDiscreteScheduler(
								    beta_schedule='''exp'''		,     num_train_timesteps=1_024		,     prediction_type='''sample'''		,     use_karras_sigmas=_snake_case		,     clip_sample=_snake_case		,     clip_sample_range=1.0		,     )
								__a										=    {
								    '''prior''': prior,
								    '''image_encoder''': image_encoder,
								    '''image_processor''': image_processor,
								    '''renderer''': renderer,
								    '''scheduler''': scheduler,
								}

								return components




				def 				SCREAMING_SNAKE_CASE_		(     self		,     _snake_case		,     _snake_case=0				)    ->      List[Any]:


								'''simple docstring'''




								__a										=    floats_tensor((1, 3, 64, 64)		,     rng=random.Random(_snake_case				)				).to(_snake_case				)

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




				def 				SCREAMING_SNAKE_CASE_		(     self				)    ->      Tuple:


								'''simple docstring'''




								__a										=    '''cpu'''

								__a										=    self.get_dummy_components()

								__a										=    self.pipeline_class(**_snake_case				)
								__a										=    pipe.to(_snake_case				)

								pipe.set_progress_bar_config(disable=_snake_case				)

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

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

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


								'''simple docstring'''




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




				def 				SCREAMING_SNAKE_CASE_		(     self				)    ->      List[str]:


								'''simple docstring'''




								__a										=    torch_device == '''cpu'''
								__a										=    True
								self._test_inference_batch_single_identical(
								    batch_size=2		,     test_max_difference=_snake_case		,     relax_max_difference=_snake_case		,     )




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


								'''simple docstring'''




								__a										=    self.get_dummy_components()
								__a										=    self.pipeline_class(**_snake_case				)
								__a										=    pipe.to(_snake_case				)
								pipe.set_progress_bar_config(disable=_snake_case				)

								__a										=    1
								__a										=    2

								__a										=    self.get_dummy_inputs(_snake_case				)

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

								__a										=    pipe(**_snake_case		,     num_images_per_prompt=_snake_case				)[0]

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







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




				def 				SCREAMING_SNAKE_CASE_		(     self				)    ->      List[str]:


								'''simple docstring'''




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




				def 				SCREAMING_SNAKE_CASE_		(     self				)    ->      str:


								'''simple docstring'''




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

								__a										=    torch.Generator(device=_snake_case				).manual_seed(0				)

								__a										=    pipe(
								    _snake_case		,     generator=_snake_case		,     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(_snake_case		,     _snake_case				)

import warnings

from ...processing_utils import ProcessorMixin
from ...tokenization_utils_base import BatchEncoding







class 			__A( a			):
				snake_case_						   =    ['''image_processor''', '''tokenizer''']
				snake_case_						   =    '''ChineseCLIPImageProcessor'''
				snake_case_						   =    ('''BertTokenizer''', '''BertTokenizerFast''')




				def __init__(     self		,     _snake_case=None		,     _snake_case=None		,     **_snake_case				)    ->      Tuple:


								'''simple docstring'''




								__a										=    None
								if "feature_extractor" in kwargs:
												warnings.warn(
												    '''The `feature_extractor` argument is deprecated and will be removed in v5, use `image_processor`'''
												    ''' instead.'''		,     _snake_case		,     )
												__a										=    kwargs.pop('''feature_extractor'''				)

								__a										=    image_processor if image_processor is not None else feature_extractor
								if image_processor is None:
												raise ValueError('''You need to specify an `image_processor`.'''				)
								if tokenizer is None:
												raise ValueError('''You need to specify a `tokenizer`.'''				)

								super().__init__(_snake_case		,     _snake_case				)
								__a										=    self.image_processor




				def __call__(     self		,     _snake_case=None		,     _snake_case=None		,     _snake_case=None		,     **_snake_case				)    ->      Optional[Any]:


								'''simple docstring'''




								if text is None and images is None:
												raise ValueError('''You have to specify either text or images. Both cannot be none.'''				)

								if text is not None:
												__a										=    self.tokenizer(_snake_case		,     return_tensors=_snake_case		,     **_snake_case				)

								if images is not None:
												__a										=    self.image_processor(_snake_case		,     return_tensors=_snake_case		,     **_snake_case				)

								if text is not None and images is not None:
												__a										=    image_features.pixel_values
												return encoding
								elif text is not None:
												return encoding
								else:
												return BatchEncoding(data=dict(**_snake_case				)		,     tensor_type=_snake_case				)




				def 				SCREAMING_SNAKE_CASE_		(     self		,     *_snake_case		,     **_snake_case				)    ->      str:


								'''simple docstring'''




								return self.tokenizer.batch_decode(*_snake_case		,     **_snake_case				)




				def 				SCREAMING_SNAKE_CASE_		(     self		,     *_snake_case		,     **_snake_case				)    ->      Dict:


								'''simple docstring'''




								return self.tokenizer.decode(*_snake_case		,     **_snake_case				)




				@property
				def 				SCREAMING_SNAKE_CASE_		(     self				)    ->      Any:


								'''simple docstring'''




								__a										=    self.tokenizer.model_input_names
								__a										=    self.image_processor.model_input_names
								return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names				)				)




				@property
				def 				SCREAMING_SNAKE_CASE_		(     self				)    ->      Optional[Any]:


								'''simple docstring'''




								warnings.warn(
								    '''`feature_extractor_class` is deprecated and will be removed in v5. Use `image_processor_class` instead.'''		,     _snake_case		,     )
								return self.image_processor_class