



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       _snake_case   (	lowercase__  ,						unittest.TestCase							):
 __lowerCAmelCase  :     List[Any]        =			ShapEPipeline
 __lowerCAmelCase  :     int        =			['prompt']
 __lowerCAmelCase  :     int        =			['prompt']
 __lowerCAmelCase  :     Optional[int]        =			[
     'num_images_per_prompt',
     'num_inference_steps',
     'generator',
     'latents',
     'guidance_scale',
     'frame_size',
     'output_type',
     'return_dict',
 ]
 __lowerCAmelCase  :     Any        =			False


 @property
 def      lowercase__     (					self):


   '''simple docstring'''




   return 32


 @property
 def      lowercase__     (					self):


   '''simple docstring'''




   return 32


 @property
 def      lowercase__     (					self):


   '''simple docstring'''




   return self.time_input_dim * 4


 @property
 def      lowercase__     (					self):


   '''simple docstring'''




   return 8


 @property
 def      lowercase__     (					self):


   '''simple docstring'''




   lowercase__  :   List[Any]    		=      CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip""")
   return tokenizer


 @property
 def      lowercase__     (					self):


   '''simple docstring'''




   torch.manual_seed(0)
   lowercase__  :   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-05					,  num_attention_heads=4					,  num_hidden_layers=5					,  pad_token_id=1					,  vocab_size=10_00					,  )
   return CLIPTextModelWithProjection(SCREAMING_SNAKE_CASE_)


 @property
 def      lowercase__     (					self):


   '''simple docstring'''




   torch.manual_seed(0)

   lowercase__  :   Union[str, 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,
   }

   lowercase__  :   Any    		=      PriorTransformer(**SCREAMING_SNAKE_CASE_)
   return model


 @property
 def      lowercase__     (					self):


   '''simple docstring'''




   torch.manual_seed(0)

   lowercase__  :   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,
       ),
   }
   lowercase__  :   Optional[int]    		=      ShapERenderer(**SCREAMING_SNAKE_CASE_)
   return model


 def      lowercase__     (					self):


   '''simple docstring'''




   lowercase__  :   List[str]    		=      self.dummy_prior
   lowercase__  :   Tuple    		=      self.dummy_text_encoder
   lowercase__  :   Optional[Any]    		=      self.dummy_tokenizer
   lowercase__  :   Optional[Any]    		=      self.dummy_renderer

   lowercase__  :   Tuple    		=      HeunDiscreteScheduler(
       beta_schedule="""exp"""					,  num_train_timesteps=10_24					,  prediction_type="""sample"""					,  use_karras_sigmas=SCREAMING_SNAKE_CASE_					,  clip_sample=SCREAMING_SNAKE_CASE_					,  clip_sample_range=1.0					,  )
   lowercase__  :   Optional[Any]    		=      {
       "prior": prior,
       "text_encoder": text_encoder,
       "tokenizer": tokenizer,
       "renderer": renderer,
       "scheduler": scheduler,
   }

   return components


 def      lowercase__     (					self					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_=0):


   '''simple docstring'''




   if str(SCREAMING_SNAKE_CASE_).startswith("""mps"""):
     lowercase__  :   Tuple    		=      torch.manual_seed(SCREAMING_SNAKE_CASE_)
   else:
     lowercase__  :   Optional[int]    		=      torch.Generator(device=SCREAMING_SNAKE_CASE_).manual_seed(SCREAMING_SNAKE_CASE_)
   lowercase__  :   List[Any]    		=      {
       "prompt": "horse",
       "generator": generator,
       "num_inference_steps": 1,
       "frame_size": 32,
       "output_type": "np",
   }
   return inputs


 def      lowercase__     (					self):


   '''simple docstring'''




   lowercase__  :   Optional[int]    		=      "cpu"

   lowercase__  :   str    		=      self.get_dummy_components()

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

   pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_)

   lowercase__  :   List[Any]    		=      pipe(**self.get_dummy_inputs(SCREAMING_SNAKE_CASE_))
   lowercase__  :   Dict    		=      output.images[0]
   lowercase__  :   str    		=      image[0, -3:, -3:, -1]

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

   lowercase__  :   Optional[Any]    		=      np.array(
       [
           0.0_0_0_3_9_2_1_6,
           0.0_0_0_3_9_2_1_6,
           0.0_0_0_3_9_2_1_6,
           0.0_0_0_3_9_2_1_6,
           0.0_0_0_3_9_2_1_6,
           0.0_0_0_3_9_2_1_6,
           0.0_0_0_3_9_2_1_6,
           0.0_0_0_3_9_2_1_6,
           0.0_0_0_3_9_2_1_6,
       ])

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


 def      lowercase__     (					self):


   '''simple docstring'''




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


 def      lowercase__     (					self):


   '''simple docstring'''




   lowercase__  :   Union[str, Any]    		=      torch_device == "cpu"
   lowercase__  :   List[str]    		=      True

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



 def      lowercase__     (					self):


   '''simple docstring'''




   lowercase__  :   Tuple    		=      self.get_dummy_components()
   lowercase__  :   List[str]    		=      self.pipeline_class(**SCREAMING_SNAKE_CASE_)
   lowercase__  :   Optional[Any]    		=      pipe.to(SCREAMING_SNAKE_CASE_)
   pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_)

   lowercase__  :   int    		=      1
   lowercase__  :   Dict    		=      2

   lowercase__  :   int    		=      self.get_dummy_inputs(SCREAMING_SNAKE_CASE_)

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

   lowercase__  :   Union[str, Any]    		=      pipe(**SCREAMING_SNAKE_CASE_					,  num_images_per_prompt=SCREAMING_SNAKE_CASE_)[0]

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



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


 def      lowercase__     (					self):


   '''simple docstring'''




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



 def      lowercase__     (					self):


   '''simple docstring'''




   lowercase__  :   Optional[Any]    		=      load_numpy(
       """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
       """/shap_e/test_shap_e_np_out.npy""")
   lowercase__  :   Tuple    		=      ShapEPipeline.from_pretrained("""openai/shap-e""")
   lowercase__  :   Union[str, Any]    		=      pipe.to(SCREAMING_SNAKE_CASE_)
   pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_)

   lowercase__  :   Tuple    		=      torch.Generator(device=SCREAMING_SNAKE_CASE_).manual_seed(0)

   lowercase__  :   Union[str, Any]    		=      pipe(
       """a shark"""					,  generator=SCREAMING_SNAKE_CASE_					,  guidance_scale=1_5.0					,  num_inference_steps=64					,  frame_size=64					,  output_type="""np"""					,  ).images[0]

   assert images.shape == (20, 64, 64, 3)

   assert_mean_pixel_difference(SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_)








import math
import random
from typing import Any

from .hill_climbing import SearchProblem

def 		UpperCamelCase	(      lowercase_ , lowercase_ = True , lowercase_ = math.inf , lowercase_ = -math.inf , lowercase_ = math.inf , lowercase_ = -math.inf , lowercase_ = False , lowercase_ = 1_00 , lowercase_ = 0.01 , lowercase_ = 1 , )		->	Any:


  '''simple docstring'''






  lowercase__  :   Any    		=      False
  lowercase__  :   Union[str, Any]    		=      search_prob
  lowercase__  :   int    		=      start_temperate
  lowercase__  :   Union[str, Any]    		=      []
  lowercase__  :   List[str]    		=      0
  lowercase__  :   Optional[int]    		=      None

  while not search_end:
    lowercase__  :   int    		=      current_state.score()
    if best_state is None or current_score > best_state.score():
      lowercase__  :   str    		=      current_state
    scores.append(lowercase_   )
    iterations += 1
    lowercase__  :   Union[str, Any]    		=      None
    lowercase__  :   int    		=      current_state.get_neighbors()
    while (
        next_state is None and neighbors
    ):  # till we do not find a neighbor that we can move to
      lowercase__  :   List[str]    		=      random.randint(0 , len(lowercase_   ) - 1   )  # picking a random neighbor
      lowercase__  :   Optional[int]    		=      neighbors.pop(lowercase_   )
      lowercase__  :   Dict    		=      picked_neighbor.score() - current_score

      if (
          picked_neighbor.x > max_x
          or picked_neighbor.x < min_x
          or picked_neighbor.y > max_y
          or picked_neighbor.y < min_y
      ):
        continue  # neighbor outside our bounds

      if not find_max:
        lowercase__  :   Any    		=      change * -1  # in case we are finding minimum
      if change > 0:  # improves the solution
        lowercase__  :   Tuple    		=      picked_neighbor
      else:
        lowercase__  :   Union[str, Any]    		=      (math.e) ** (
            change / current_temp
        )  # probability generation function
        if random.random() < probability:  # random number within probability
          lowercase__  :   Optional[int]    		=      picked_neighbor
    lowercase__  :   Dict    		=      current_temp - (current_temp * rate_of_decrease)

    if current_temp < threshold_temp or next_state is None:
      # temperature below threshold, or could not find a suitable neighbor
      lowercase__  :   Optional[Any]    		=      True
    else:
      lowercase__  :   List[Any]    		=      next_state

  if visualization:
    from matplotlib import pyplot as plt

    plt.plot(range(lowercase_   ) , lowercase_   )
    plt.xlabel("""Iterations"""   )
    plt.ylabel("""Function values"""   )
    plt.show()
  return best_state


if __name__ == "__main__":

     def 		UpperCamelCase	(      lowercase_ , lowercase_   )		->	str:


       '''simple docstring'''






       return (x**2) + (y**2)

     # starting the problem with initial coordinates (12, 47)
     lowerCamelCase__    :		str          =   SearchProblem(x=1_2, y=4_7, step_size=1, function_to_optimize=test_fa)
     lowerCamelCase__    :		List[Any]          =   simulated_annealing(
         prob, find_max=False, max_x=1_0_0, min_x=5, max_y=5_0, min_y=-5, visualization=True
     )
     print(
         """The minimum score for f(x, y) = x^2 + y^2 with the domain 100 > x > 5 """
         f'''and 50 > y > - 5 found via hill climbing: {local_min.score()}'''
     )

     # starting the problem with initial coordinates (12, 47)
     lowerCamelCase__    :		List[Any]          =   SearchProblem(x=1_2, y=4_7, step_size=1, function_to_optimize=test_fa)
     lowerCamelCase__    :		Optional[int]          =   simulated_annealing(
         prob, find_max=True, max_x=1_0_0, min_x=5, max_y=5_0, min_y=-5, visualization=True
     )
     print(
         """The maximum score for f(x, y) = x^2 + y^2 with the domain 100 > x > 5 """
         f'''and 50 > y > - 5 found via hill climbing: {local_min.score()}'''
     )

     def 		UpperCamelCase	(      lowercase_ , lowercase_   )		->	Union[str, Any]:


       '''simple docstring'''






       return (3 * x**2) - (6 * y)

     lowerCamelCase__    :		Tuple          =   SearchProblem(x=3, y=4, step_size=1, function_to_optimize=test_fa)
     lowerCamelCase__    :		Union[str, Any]          =   simulated_annealing(prob, find_max=False, visualization=True)
     print(
         """The minimum score for f(x, y) = 3*x^2 - 6*y found via hill climbing: """
         f'''{local_min.score()}'''
     )

     lowerCamelCase__    :		Dict          =   SearchProblem(x=3, y=4, step_size=1, function_to_optimize=test_fa)
     lowerCamelCase__    :		Tuple          =   simulated_annealing(prob, find_max=True, visualization=True)



     print(
         """The maximum score for f(x, y) = 3*x^2 - 6*y found via hill climbing: """
         f'''{local_min.score()}'''
     )






