





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







  '''simple docstring'''


  SCREAMING_SNAKE_CASE__				:Union[str, Any]						  =       ShapEPipeline
  SCREAMING_SNAKE_CASE__				:str						  =       ["prompt"]
  SCREAMING_SNAKE_CASE__				:List[str]						  =       ["prompt"]
  SCREAMING_SNAKE_CASE__				:Dict						  =       [
      "num_images_per_prompt",
      "num_inference_steps",
      "generator",
      "latents",
      "guidance_scale",
      "frame_size",
      "output_type",
      "return_dict",
  ]
  SCREAMING_SNAKE_CASE__				:int						  =       False

  @property
  def  __SCREAMING_SNAKE_CASE					(		self :							Dict						)		->					int:
        return 32

  @property
  def  __SCREAMING_SNAKE_CASE					(		self :							List[Any]						)		->					int:
        return 32

  @property
  def  __SCREAMING_SNAKE_CASE					(		self :							Dict						)		->					Optional[int]:
        return self.time_input_dim * 4

  @property
  def  __SCREAMING_SNAKE_CASE					(		self :							List[Any]						)		->					int:
        return 8

  @property
  def  __SCREAMING_SNAKE_CASE					(		self :							Optional[Any]						)		->					List[str]:
        _UpperCamelCase   :						List[Any]    =							CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"						)
        return tokenizer

  @property
  def  __SCREAMING_SNAKE_CASE					(		self :							List[str]						)		->					List[Any]:
        torch.manual_seed(0						)
        _UpperCamelCase   :						str    =							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=1000       , )
        return CLIPTextModelWithProjection(__a						)

  @property
  def  __SCREAMING_SNAKE_CASE					(		self :							Tuple						)		->					Tuple:
        torch.manual_seed(0						)

        _UpperCamelCase   :						Tuple    =							{
            "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   :						Union[str, Any]    =							PriorTransformer(**__a						)
        return model

  @property
  def  __SCREAMING_SNAKE_CASE					(		self :							int						)		->					Any:
        torch.manual_seed(0						)

        _UpperCamelCase   :						Optional[int]    =							{
            "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   :						List[Any]    =							ShapERenderer(**__a						)
        return model

  def  __SCREAMING_SNAKE_CASE					(		self :							Dict						)		->					Union[str, Any]:
        _UpperCamelCase   :						Optional[Any]    =							self.dummy_prior
        _UpperCamelCase   :						List[Any]    =							self.dummy_text_encoder
        _UpperCamelCase   :						Any    =							self.dummy_tokenizer
        _UpperCamelCase   :						str    =							self.dummy_renderer

        _UpperCamelCase   :						str    =							HeunDiscreteScheduler(
            beta_schedule="exp"       , num_train_timesteps=1024       , prediction_type="sample"       , use_karras_sigmas=__a       , clip_sample=__a       , clip_sample_range=1.0       , )
        _UpperCamelCase   :						List[Any]    =							{
            "prior": prior,
            "text_encoder": text_encoder,
            "tokenizer": tokenizer,
            "renderer": renderer,
            "scheduler": scheduler,
        }

        return components

  def  __SCREAMING_SNAKE_CASE					(		self :							List[str]       , __a :							Tuple       , __a :							Dict=0						)		->					Any:
        if str(__a						).startswith("mps"						):
              _UpperCamelCase   :						Optional[int]    =							torch.manual_seed(__a						)
        else:
              _UpperCamelCase   :						Dict    =							torch.Generator(device=__a						).manual_seed(__a						)
        _UpperCamelCase   :						Any    =							{
            "prompt": "horse",
            "generator": generator,
            "num_inference_steps": 1,
            "frame_size": 32,
            "output_type": "np",
        }
        return inputs

  def  __SCREAMING_SNAKE_CASE					(		self :							Tuple						)		->					Optional[Any]:
        _UpperCamelCase   :						Union[str, Any]    =							"cpu"

        _UpperCamelCase   :						Union[str, Any]    =							self.get_dummy_components()

        _UpperCamelCase   :						int    =							self.pipeline_class(**__a						)
        _UpperCamelCase   :						Dict    =							pipe.to(__a						)

        pipe.set_progress_bar_config(disable=__a						)

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

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

        _UpperCamelCase   :						Optional[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 :							List[str]						)		->					Tuple:
        # NOTE: Larger batch sizes cause this test to timeout, only test on smaller batches
        self._test_inference_batch_consistent(batch_sizes=[1, 2]						)

  def  __SCREAMING_SNAKE_CASE					(		self :							int						)		->					Union[str, Any]:
        _UpperCamelCase   :						str    =							torch_device == "cpu"
        _UpperCamelCase   :						Optional[Any]    =							True

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


  def  __SCREAMING_SNAKE_CASE					(		self :							List[str]						)		->					Any:
        _UpperCamelCase   :						str    =							self.get_dummy_components()
        _UpperCamelCase   :						Tuple    =							self.pipeline_class(**__a						)
        _UpperCamelCase   :						Union[str, Any]    =							pipe.to(__a						)
        pipe.set_progress_bar_config(disable=__a						)

        _UpperCamelCase   :						Dict    =							1
        _UpperCamelCase   :						List[str]    =							2

        _UpperCamelCase   :						Dict    =							self.get_dummy_inputs(__a						)

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

        _UpperCamelCase   :						List[str]    =							pipe(**__a       , num_images_per_prompt=__a						)[0]

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





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







  '''simple docstring'''


  def  __SCREAMING_SNAKE_CASE					(		self :							Any						)		->					Tuple:
        # clean up the VRAM after each test
        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()


  def  __SCREAMING_SNAKE_CASE					(		self :							List[Any]						)		->					Tuple:
        _UpperCamelCase   :						List[Any]    =							load_numpy(
            "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
            "/shap_e/test_shap_e_np_out.npy"						)
        _UpperCamelCase   :						Dict    =							ShapEPipeline.from_pretrained("openai/shap-e"						)
        _UpperCamelCase   :						str    =							pipe.to(__a						)
        pipe.set_progress_bar_config(disable=__a						)

        _UpperCamelCase   :						int    =							torch.Generator(device=__a						).manual_seed(0						)

        _UpperCamelCase   :						List[str]    =							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 collections.abc import Callable

import numpy as np




def        lowercase__     (       lowercase_		,lowercase_		,lowercase_		,lowercase_		,lowercase_      )       ->		np.ndarray:




      """simple docstring"""

      _UpperCamelCase   :						List[str]    =							int(np.ceil((x_end - xa) / step_size      )      )
      _UpperCamelCase   :						Tuple    =							np.zeros((n + 1,)      )
      _UpperCamelCase   :						List[str]    =							ya
      _UpperCamelCase   :						List[Any]    =							xa

      for k in range(lowercase_      ):
            _UpperCamelCase   :						List[str]    =							y[k] + step_size * ode_func(lowercase_		,y[k]      )
            x += step_size

      return y


if __name__ == "__main__":
  import doctest

  doctest.testmod()
