




"""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 					__magic_name__							( UpperCAmelCase__		,      unittest.TestCase       ):






				'''simple docstring'''
				__UpperCamelCase    						=		ShapEPipeline
				__UpperCamelCase    						=		["prompt"]
				__UpperCamelCase    						=		["prompt"]
				__UpperCamelCase    						=		[
				    "num_images_per_prompt",
				    "num_inference_steps",
				    "generator",
				    "latents",
				    "guidance_scale",
				    "frame_size",
				    "output_type",
				    "return_dict",
				]
				__UpperCamelCase    						=		False





				@property
				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							return 32





				@property
				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							return 32





				@property
				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							return self.time_input_dim * 4





				@property
				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							return 8





				@property
				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							lowerCamelCase       =      CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""				)
							return tokenizer





				@property
				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							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-0_5	,				num_attention_heads=4	,				num_hidden_layers=5	,				pad_token_id=1	,				vocab_size=1_000	,				)
							return CLIPTextModelWithProjection(_a				)





				@property
				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							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(**_a				)
							return model





				@property
				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							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(**_a				)
							return model





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							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=1_024	,				prediction_type="""sample"""	,				use_karras_sigmas=_a	,				clip_sample=_a	,				clip_sample_range=1.0	,				)
							lowerCamelCase       =      {
							    """prior""": prior,
							    """text_encoder""": text_encoder,
							    """tokenizer""": tokenizer,
							    """renderer""": renderer,
							    """scheduler""": scheduler,
							}

							return components





				def      _lowerCAmelCase       ( self	,				_a	,				_a=0				):





							"""simple docstring"""
							if str(_a				).startswith("""mps"""				):
										lowerCamelCase       =      torch.manual_seed(_a				)
							else:
										lowerCamelCase       =      torch.Generator(device=_a				).manual_seed(_a				)
							lowerCamelCase       =      {
							    """prompt""": """horse""",
							    """generator""": generator,
							    """num_inference_steps""": 1,
							    """frame_size""": 32,
							    """output_type""": """np""",
							}
							return inputs





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							lowerCamelCase       =      """cpu"""

							lowerCamelCase       =      self.get_dummy_components()

							lowerCamelCase       =      self.pipeline_class(**_a				)
							lowerCamelCase       =      pipe.to(_a				)

							pipe.set_progress_bar_config(disable=_a				)

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

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

							lowerCamelCase       =      np.array(
							    [
							        0.00_039_216,
							        0.00_039_216,
							        0.00_039_216,
							        0.00_039_216,
							        0.00_039_216,
							        0.00_039_216,
							        0.00_039_216,
							        0.00_039_216,
							        0.00_039_216,
							    ]				)

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





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							# 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				):





							"""simple docstring"""
							lowerCamelCase       =      torch_device == """cpu"""
							lowerCamelCase       =      True

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





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							lowerCamelCase       =      self.get_dummy_components()
							lowerCamelCase       =      self.pipeline_class(**_a				)
							lowerCamelCase       =      pipe.to(_a				)
							pipe.set_progress_bar_config(disable=_a				)

							lowerCamelCase       =      1
							lowerCamelCase       =      2

							lowerCamelCase       =      self.get_dummy_inputs(_a				)

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

							lowerCamelCase       =      pipe(**_a	,				num_images_per_prompt=_a				)[0]

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






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






				'''simple docstring'''





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							# clean up the VRAM after each test
							super().tearDown()
							gc.collect()
							torch.cuda.empty_cache()





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							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(_a				)
							pipe.set_progress_bar_config(disable=_a				)

							lowerCamelCase       =      torch.Generator(device=_a				).manual_seed(0				)

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









"""simple docstring"""




from typing import TYPE_CHECKING

# rely on isort to merge the imports
from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


lowerCAmelCase							:    List[str]  							=							{
    """configuration_autoformer""": [
        """AUTOFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP""",
        """AutoformerConfig""",
    ],
}

try:
					if not is_torch_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					lowerCAmelCase							:    Tuple  							=							[
					    """AUTOFORMER_PRETRAINED_MODEL_ARCHIVE_LIST""",
					    """AutoformerForPrediction""",
					    """AutoformerModel""",
					    """AutoformerPreTrainedModel""",
					]


if TYPE_CHECKING:
					from .configuration_autoformer import (
					    AUTOFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP,
					    AutoformerConfig,
					)

					try:
										if not is_torch_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .modeling_autoformer import (
										    AUTOFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
										    AutoformerForPrediction,
										    AutoformerModel,
										    AutoformerPreTrainedModel,
										)

else:
					import sys

					lowerCAmelCase							:    str  							=							_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)



