


'''simple docstring'''



from dataclasses import dataclass
from typing import List, Optional, Union

import numpy as np
import PIL

from ...utils import BaseOutput, OptionalDependencyNotAvailable, is_torch_available, is_transformers_available
from .timesteps import (
    fastaa_timesteps,
    smartaa_timesteps,
    smartaa_timesteps,
    smartaaa_timesteps,
    smartaaa_timesteps,
    superaa_timesteps,
    superaa_timesteps,
    superaaa_timesteps,
)






@dataclass
class 	UpperCamelCase__(       lowerCAmelCase	):
				__magic_name__    :				Union[List[PIL.Image.Image], np.ndarray]
				__magic_name__    :				Optional[List[bool]]
				__magic_name__    :				Optional[List[bool]]


try:
	if not (is_transformers_available() and is_torch_available()):
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	from ...utils.dummy_torch_and_transformers_objects import *  # noqa F403
else:
	from .pipeline_if import IFPipeline
	from .pipeline_if_imgaimg import IFImgaImgPipeline
	from .pipeline_if_imgaimg_superresolution import IFImgaImgSuperResolutionPipeline
	from .pipeline_if_inpainting import IFInpaintingPipeline
	from .pipeline_if_inpainting_superresolution import IFInpaintingSuperResolutionPipeline
	from .pipeline_if_superresolution import IFSuperResolutionPipeline
	from .safety_checker import IFSafetyChecker
	from .watermark import IFWatermarker









'''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__(       lowerCAmelCase						,    unittest.TestCase	):
				__magic_name__    :				Any  		=	ShapEPipeline
				__magic_name__    :				Tuple  		=	["prompt"]
				__magic_name__    :				Optional[int]  		=	["prompt"]
				__magic_name__    :				Dict  		=	[
				    "num_images_per_prompt",
				    "num_inference_steps",
				    "generator",
				    "latents",
				    "guidance_scale",
				    "frame_size",
				    "output_type",
				    "return_dict",
				]
				__magic_name__    :				Optional[int]  		=	False




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

										"""simple docstring"""




										return 32




				@property
				def 					a__(		self      :			Dict				)->					Dict:

										"""simple docstring"""




										return 32




				@property
				def 					a__(		self      :			Optional[Any]				)->					List[str]:

										"""simple docstring"""




										return self.time_input_dim * 4




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

										"""simple docstring"""




										return 8




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

										"""simple docstring"""




										UpperCAmelCase					    =     CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''				)
										return tokenizer




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

										"""simple docstring"""




										torch.manual_seed(0				)
										UpperCAmelCase					    =     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=1000				,   )
										return CLIPTextModelWithProjection(lowerCAmelCase				)




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

										"""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''',
										    '''encoder_hid_proj_type''': None,
										    '''added_emb_type''': None,
										}

										UpperCAmelCase					    =     PriorTransformer(**lowerCAmelCase				)
										return model




				@property
				def 					a__(		self      :			List[Any]				)->					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(**lowerCAmelCase				)
										return model




				def 					a__(		self      :			Any				)->					Tuple:

										"""simple docstring"""




										UpperCAmelCase					    =     self.dummy_prior
										UpperCAmelCase					    =     self.dummy_text_encoder
										UpperCAmelCase					    =     self.dummy_tokenizer
										UpperCAmelCase					    =     self.dummy_renderer

										UpperCAmelCase					    =     HeunDiscreteScheduler(
										    beta_schedule='''exp'''				,   num_train_timesteps=1024				,   prediction_type='''sample'''				,   use_karras_sigmas=lowerCAmelCase				,   clip_sample=lowerCAmelCase				,   clip_sample_range=1.0				,   )
										UpperCAmelCase					    =     {
										    '''prior''': prior,
										    '''text_encoder''': text_encoder,
										    '''tokenizer''': tokenizer,
										    '''renderer''': renderer,
										    '''scheduler''': scheduler,
										}

										return components




				def 					a__(		self      :			Tuple				,   lowerCAmelCase      :			Union[str, Any]				,   lowerCAmelCase      :			List[Any]=0				)->					Optional[Any]:

										"""simple docstring"""




										if str(lowerCAmelCase				).startswith('''mps'''				):
																UpperCAmelCase					    =     torch.manual_seed(lowerCAmelCase				)
										else:
																UpperCAmelCase					    =     torch.Generator(device=lowerCAmelCase				).manual_seed(lowerCAmelCase				)
										UpperCAmelCase					    =     {
										    '''prompt''': '''horse''',
										    '''generator''': generator,
										    '''num_inference_steps''': 1,
										    '''frame_size''': 32,
										    '''output_type''': '''np''',
										}
										return inputs




				def 					a__(		self      :			int				)->					Dict:

										"""simple docstring"""




										UpperCAmelCase					    =     '''cpu'''

										UpperCAmelCase					    =     self.get_dummy_components()

										UpperCAmelCase					    =     self.pipeline_class(**lowerCAmelCase				)
										UpperCAmelCase					    =     pipe.to(lowerCAmelCase				)

										pipe.set_progress_bar_config(disable=lowerCAmelCase				)

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

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

										UpperCAmelCase					    =     np.array(
										    [
										        0.00039216,
										        0.00039216,
										        0.00039216,
										        0.00039216,
										        0.00039216,
										        0.00039216,
										        0.00039216,
										        0.00039216,
										        0.00039216,
										    ]				)

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




				def 					a__(		self      :			Union[str, Any]				)->					Optional[int]:

										"""simple docstring"""




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




				def 					a__(		self      :			Optional[int]				)->					str:

										"""simple docstring"""




										UpperCAmelCase					    =     torch_device == '''cpu'''
										UpperCAmelCase					    =     True

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





				def 					a__(		self      :			int				)->					List[str]:

										"""simple docstring"""




										UpperCAmelCase					    =     self.get_dummy_components()
										UpperCAmelCase					    =     self.pipeline_class(**lowerCAmelCase				)
										UpperCAmelCase					    =     pipe.to(lowerCAmelCase				)
										pipe.set_progress_bar_config(disable=lowerCAmelCase				)

										UpperCAmelCase					    =     1
										UpperCAmelCase					    =     2

										UpperCAmelCase					    =     self.get_dummy_inputs(lowerCAmelCase				)

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

										UpperCAmelCase					    =     pipe(**lowerCAmelCase				,   num_images_per_prompt=lowerCAmelCase				)[0]

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







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




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

										"""simple docstring"""




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





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

										"""simple docstring"""




										UpperCAmelCase					    =     load_numpy(
										    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
										    '''/shap_e/test_shap_e_np_out.npy'''				)
										UpperCAmelCase					    =     ShapEPipeline.from_pretrained('''openai/shap-e'''				)
										UpperCAmelCase					    =     pipe.to(lowerCAmelCase				)
										pipe.set_progress_bar_config(disable=lowerCAmelCase				)

										UpperCAmelCase					    =     torch.Generator(device=lowerCAmelCase				).manual_seed(0				)

										UpperCAmelCase					    =     pipe(
										    '''a shark'''				,   generator=lowerCAmelCase				,   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(lowerCAmelCase				,   lowerCAmelCase				)





