






"""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 							lowerCAmelCase__				(			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 _SCREAMING_SNAKE_CASE (						self     :  int):




          '''simple docstring'''


          return 32
      @property
      def _SCREAMING_SNAKE_CASE (						self     :  Tuple):




          '''simple docstring'''


          return 32
      @property
      def _SCREAMING_SNAKE_CASE (						self     :  Tuple):




          '''simple docstring'''


          return self.time_input_dim * 4
      @property
      def _SCREAMING_SNAKE_CASE (						self     :  int):




          '''simple docstring'''


          return 8
      @property
      def _SCREAMING_SNAKE_CASE (						self     :  Optional[Any]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  List[Any]    						=		CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip''')
          return tokenizer
      @property
      def _SCREAMING_SNAKE_CASE (						self     :  List[Any]):




          '''simple docstring'''


          torch.manual_seed(0)
          SCREAMING_SNAKE_CASE_      :  List[Any]    						=		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(lowercase_)
      @property
      def _SCREAMING_SNAKE_CASE (						self     :  Union[str, Any]):




          '''simple docstring'''


          torch.manual_seed(0)

          SCREAMING_SNAKE_CASE_      :  Optional[Any]    						=		{
              '''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,
          }

          SCREAMING_SNAKE_CASE_      :  Optional[Any]    						=		PriorTransformer(**lowercase_)
          return model
      @property
      def _SCREAMING_SNAKE_CASE (						self     :  Union[str, Any]):




          '''simple docstring'''


          torch.manual_seed(0)

          SCREAMING_SNAKE_CASE_      :  List[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_      :  Optional[Any]    						=		ShapERenderer(**lowercase_)
          return model
      def _SCREAMING_SNAKE_CASE (						self     :  Optional[int]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  Dict    						=		self.dummy_prior
          SCREAMING_SNAKE_CASE_      :  Union[str, Any]    						=		self.dummy_text_encoder
          SCREAMING_SNAKE_CASE_      :  Any    						=		self.dummy_tokenizer
          SCREAMING_SNAKE_CASE_      :  int    						=		self.dummy_renderer

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

          return components
      def _SCREAMING_SNAKE_CASE (						self     :  Any						,	lowercase_     :  str						,	lowercase_     :  int=0):




          '''simple docstring'''


          if str(lowercase_).startswith('''mps'''):
              SCREAMING_SNAKE_CASE_      :  str    						=		torch.manual_seed(lowercase_)
          else:
              SCREAMING_SNAKE_CASE_      :  Any    						=		torch.Generator(device=lowercase_).manual_seed(lowercase_)
          SCREAMING_SNAKE_CASE_      :  Any    						=		{
              '''prompt''': '''horse''',
              '''generator''': generator,
              '''num_inference_steps''': 1,
              '''frame_size''': 32,
              '''output_type''': '''np''',
          }
          return inputs
      def _SCREAMING_SNAKE_CASE (						self     :  Optional[Any]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  str    						=		'''cpu'''

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

          SCREAMING_SNAKE_CASE_      :  List[str]    						=		self.pipeline_class(**lowercase_)
          SCREAMING_SNAKE_CASE_      :  Optional[Any]    						=		pipe.to(lowercase_)

          pipe.set_progress_bar_config(disable=lowercase_)

          SCREAMING_SNAKE_CASE_      :  Dict    						=		pipe(**self.get_dummy_inputs(lowercase_))
          SCREAMING_SNAKE_CASE_      :  List[str]    						=		output.images[0]
          SCREAMING_SNAKE_CASE_      :  str    						=		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 _SCREAMING_SNAKE_CASE (						self     :  int):




          '''simple docstring'''


          self._test_inference_batch_consistent(batch_sizes=[1, 2])
      def _SCREAMING_SNAKE_CASE (						self     :  int):




          '''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=lowercase_						,	relax_max_difference=lowercase_						,	)



      def _SCREAMING_SNAKE_CASE (						self     :  List[str]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  int    						=		self.get_dummy_components()
          SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		self.pipeline_class(**lowercase_)
          SCREAMING_SNAKE_CASE_      :  Tuple    						=		pipe.to(lowercase_)
          pipe.set_progress_bar_config(disable=lowercase_)

          SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		1
          SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		2

          SCREAMING_SNAKE_CASE_      :  List[str]    						=		self.get_dummy_inputs(lowercase_)

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

          SCREAMING_SNAKE_CASE_      :  str    						=		pipe(**lowercase_						,	num_images_per_prompt=lowercase_)[0]

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




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






      '''simple docstring'''

      def _SCREAMING_SNAKE_CASE (						self     :  Any):




          '''simple docstring'''


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



      def _SCREAMING_SNAKE_CASE (						self     :  List[Any]):




          '''simple docstring'''


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

          SCREAMING_SNAKE_CASE_      :  Any    						=		torch.Generator(device=lowercase_).manual_seed(0)

          SCREAMING_SNAKE_CASE_      :  Any    						=		pipe(
              '''a shark'''						,	generator=lowercase_						,	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(lowercase_						,	lowercase_)










"""simple docstring"""






from . import (
    albert,
    align,
    altclip,
    audio_spectrogram_transformer,
    auto,
    autoformer,
    bark,
    bart,
    barthez,
    bartpho,
    beit,
    bert,
    bert_generation,
    bert_japanese,
    bertweet,
    big_bird,
    bigbird_pegasus,
    biogpt,
    bit,
    blenderbot,
    blenderbot_small,
    blip,
    blip_a,
    bloom,
    bridgetower,
    byta,
    camembert,
    canine,
    chinese_clip,
    clap,
    clip,
    clipseg,
    codegen,
    conditional_detr,
    convbert,
    convnext,
    convnextva,
    cpm,
    cpmant,
    ctrl,
    cvt,
    dataavec,
    deberta,
    deberta_va,
    decision_transformer,
    deformable_detr,
    deit,
    deprecated,
    deta,
    detr,
    dialogpt,
    dinat,
    distilbert,
    dit,
    donut,
    dpr,
    dpt,
    efficientformer,
    efficientnet,
    electra,
    encodec,
    encoder_decoder,
    ernie,
    ernie_m,
    esm,
    falcon,
    flaubert,
    flava,
    fnet,
    focalnet,
    fsmt,
    funnel,
    git,
    glpn,
    gpta,
    gpt_bigcode,
    gpt_neo,
    gpt_neox,
    gpt_neox_japanese,
    gpt_swa,
    gptj,
    gptsan_japanese,
    graphormer,
    groupvit,
    herbert,
    hubert,
    ibert,
    imagegpt,
    informer,
    instructblip,
    jukebox,
    layoutlm,
    layoutlmva,
    layoutlmva,
    layoutxlm,
    led,
    levit,
    lilt,
    llama,
    longformer,
    longta,
    luke,
    lxmert,
    mam_aaa,
    marian,
    markuplm,
    maskaformer,
    maskformer,
    mbart,
    mbartaa,
    mega,
    megatron_bert,
    megatron_gpta,
    mgp_str,
    mluke,
    mobilebert,
    mobilenet_va,
    mobilenet_va,
    mobilevit,
    mobilevitva,
    mpnet,
    mra,
    mta,
    musicgen,
    mvp,
    nat,
    nezha,
    nllb,
    nllb_moe,
    nystromformer,
    oneformer,
    open_llama,
    openai,
    opt,
    owlvit,
    pegasus,
    pegasus_x,
    perceiver,
    phobert,
    pixastruct,
    plbart,
    poolformer,
    prophetnet,
    qdqbert,
    rag,
    realm,
    reformer,
    regnet,
    rembert,
    resnet,
    roberta,
    roberta_prelayernorm,
    roc_bert,
    roformer,
    rwkv,
    sam,
    segformer,
    sew,
    sew_d,
    speech_encoder_decoder,
    speech_to_text,
    speech_to_text_a,
    speechta,
    splinter,
    squeezebert,
    swiftformer,
    swin,
    swinasr,
    swinva,
    switch_transformers,
    ta,
    table_transformer,
    tapas,
    time_series_transformer,
    timesformer,
    timm_backbone,
    transfo_xl,
    trocr,
    tvlt,
    umta,
    unispeech,
    unispeech_sat,
    upernet,
    videomae,
    vilt,
    vision_encoder_decoder,
    vision_text_dual_encoder,
    visual_bert,
    vit,
    vit_hybrid,
    vit_mae,
    vit_msn,
    vivit,
    wavaveca,
    wavaveca_conformer,
    wavaveca_phoneme,
    wavaveca_with_lm,
    wavlm,
    whisper,
    x_clip,
    xglm,
    xlm,
    xlm_prophetnet,
    xlm_roberta,
    xlm_roberta_xl,
    xlnet,
    xmod,
    yolos,
    yoso,
)


