
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__  (	__lowercase   ,unittest.TestCase   ):
			_SCREAMING_SNAKE_CASE				:							Any     	=				ShapEPipeline
			_SCREAMING_SNAKE_CASE				:							Any     	=				["prompt"]
			_SCREAMING_SNAKE_CASE				:							List[Any]     	=				["prompt"]
			_SCREAMING_SNAKE_CASE				:							List[str]     	=				[
			    "num_images_per_prompt",
			    "num_inference_steps",
			    "generator",
			    "latents",
			    "guidance_scale",
			    "frame_size",
			    "output_type",
			    "return_dict",
			]
			_SCREAMING_SNAKE_CASE				:							Dict     	=				False
			@property
			def 				lowerCAmelCase      (self			: Optional[Any]       ):
				return 3_2
			@property
			def 				lowerCAmelCase      (self			: Dict       ):
				return 3_2
			@property
			def 				lowerCAmelCase      (self			: Union[str, Any]       ):
				return self.time_input_dim * 4
			@property
			def 				lowerCAmelCase      (self			: Dict       ):
				return 8
			@property
			def 				lowerCAmelCase      (self			: List[str]       ):
				__a			:  Tuple		    =  CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''       )
				return tokenizer
			@property
			def 				lowerCAmelCase      (self			: Dict       ):
				torch.manual_seed(0       )
				__a			:  Tuple		    =  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=3_7 ,    layer_norm_eps=1E-05 ,    num_attention_heads=4 ,    num_hidden_layers=5 ,    pad_token_id=1 ,    vocab_size=1_0_0_0 ,    )
				return CLIPTextModelWithProjection(snake_case_       )
			@property
			def 				lowerCAmelCase      (self			: List[Any]       ):
				torch.manual_seed(0       )

				__a			:  Union[str, Any]		    =  {
				    '''num_attention_heads''': 2,
				    '''attention_head_dim''': 1_6,
				    '''embedding_dim''': self.time_input_dim,
				    '''num_embeddings''': 3_2,
				    '''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,
				}

				__a			:  Union[str, Any]		    =  PriorTransformer(**snake_case_       )
				return model
			@property
			def 				lowerCAmelCase      (self			: Tuple       ):
				torch.manual_seed(0       )

				__a			:  Optional[int]		    =  {
				    '''param_shapes''': (
				        (self.renderer_dim, 9_3),
				        (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''': 1_2,
				    '''background''': (
				        0.1,
				        0.1,
				        0.1,
				    ),
				}
				__a			:  int		    =  ShapERenderer(**snake_case_       )
				return model
			def 				lowerCAmelCase      (self			: Any       ):
				__a			:  Optional[Any]		    =  self.dummy_prior
				__a			:  Dict		    =  self.dummy_text_encoder
				__a			:  List[Any]		    =  self.dummy_tokenizer
				__a			:  Any		    =  self.dummy_renderer

				__a			:  Dict		    =  HeunDiscreteScheduler(
				    beta_schedule='''exp''' ,    num_train_timesteps=1_0_2_4 ,    prediction_type='''sample''' ,    use_karras_sigmas=snake_case_ ,    clip_sample=snake_case_ ,    clip_sample_range=1.0 ,    )
				__a			:  Union[str, Any]		    =  {
				    '''prior''': prior,
				    '''text_encoder''': text_encoder,
				    '''tokenizer''': tokenizer,
				    '''renderer''': renderer,
				    '''scheduler''': scheduler,
				}

				return components
			def 				lowerCAmelCase      (self			: str ,    snake_case_			: Dict ,    snake_case_			: Optional[int]=0       ):
				if str(snake_case_       ).startswith('''mps'''       ):
					__a			:  List[str]		    =  torch.manual_seed(snake_case_       )
				else:
					__a			:  str		    =  torch.Generator(device=snake_case_       ).manual_seed(snake_case_       )
				__a			:  Dict		    =  {
				    '''prompt''': '''horse''',
				    '''generator''': generator,
				    '''num_inference_steps''': 1,
				    '''frame_size''': 3_2,
				    '''output_type''': '''np''',
				}
				return inputs
			def 				lowerCAmelCase      (self			: int       ):
				__a			:  List[Any]		    =  '''cpu'''

				__a			:  Dict		    =  self.get_dummy_components()

				__a			:  str		    =  self.pipeline_class(**snake_case_       )
				__a			:  Optional[int]		    =  pipe.to(snake_case_       )

				pipe.set_progress_bar_config(disable=snake_case_       )

				__a			:  Any		    =  pipe(**self.get_dummy_inputs(snake_case_       )       )
				__a			:  List[str]		    =  output.images[0]
				__a			:  int		    =  image[0, -3:, -3:, -1]

				assert image.shape == (2_0, 3_2, 3_2, 3)

				__a			:  Union[str, Any]		    =  np.array(
				    [
				        0.0003_9216,
				        0.0003_9216,
				        0.0003_9216,
				        0.0003_9216,
				        0.0003_9216,
				        0.0003_9216,
				        0.0003_9216,
				        0.0003_9216,
				        0.0003_9216,
				    ]       )

				assert np.abs(image_slice.flatten() - expected_slice       ).max() < 1E-2
			def 				lowerCAmelCase      (self			: Any       ):
				# 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			: Tuple       ):
				__a			:  str		    =  torch_device == '''cpu'''
				__a			:  int		    =  True

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







			def 				lowerCAmelCase      (self			: Tuple       ):
				__a			:  List[str]		    =  self.get_dummy_components()
				__a			:  Any		    =  self.pipeline_class(**snake_case_       )
				__a			:  Optional[int]		    =  pipe.to(snake_case_       )
				pipe.set_progress_bar_config(disable=snake_case_       )

				__a			:  int		    =  1
				__a			:  List[Any]		    =  2

				__a			:  Any		    =  self.get_dummy_inputs(snake_case_       )

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

				__a			:  List[Any]		    =  pipe(**snake_case_ ,    num_images_per_prompt=snake_case_       )[0]

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





@slow
@require_torch_gpu
class 		UpperCamelCase__  (	unittest.TestCase   ):
			def 				lowerCAmelCase      (self			: Dict       ):
				# clean up the VRAM after each test
				super().tearDown()
				gc.collect()
				torch.cuda.empty_cache()







			def 				lowerCAmelCase      (self			: str       ):
				__a			:  int		    =  load_numpy(
				    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
				    '''/shap_e/test_shap_e_np_out.npy'''       )
				__a			:  Tuple		    =  ShapEPipeline.from_pretrained('''openai/shap-e'''       )
				__a			:  Union[str, Any]		    =  pipe.to(snake_case_       )
				pipe.set_progress_bar_config(disable=snake_case_       )

				__a			:  str		    =  torch.Generator(device=snake_case_       ).manual_seed(0       )

				__a			:  List[str]		    =  pipe(
				    '''a shark''' ,    generator=snake_case_ ,    guidance_scale=15.0 ,    num_inference_steps=6_4 ,    frame_size=6_4 ,    output_type='''np''' ,    ).images[0]

				assert images.shape == (2_0, 6_4, 6_4, 3)

				assert_mean_pixel_difference(snake_case_ ,    snake_case_       )


from __future__ import annotations

# This is the precision for this function which can be altered.
# It is recommended for users to keep this number greater than or equal to 10.
lowercase__          =10




def        __UpperCamelCase					(    lowerCAmelCase__							:    int  ,       lowerCAmelCase__							:    int  ,       lowerCAmelCase__							:    list[int]  ,       lowerCAmelCase__							:    int  ):
 for i in range(lowerCAmelCase__  ,       lowerCAmelCase__  ):
  if array[i] == target:
   return i
 return -1




def        __UpperCamelCase					(    lowerCAmelCase__							:    list[int]  ,       lowerCAmelCase__							:    int  ):
 __a			:  List[Any]		    =  0
 __a			:  Union[str, Any]		    =  len(lowerCAmelCase__  )
 while left <= right:
  if right - left < precision:
   return lin_search(lowerCAmelCase__  ,       lowerCAmelCase__  ,       lowerCAmelCase__  ,       lowerCAmelCase__  )

  __a			:  str		    =  (left + right) // 3 + 1
  __a			:  List[Any]		    =  2 * (left + right) // 3 + 1

  if array[one_third] == target:
   return one_third
  elif array[two_third] == target:
   return two_third

  elif target < array[one_third]:
   __a			:  List[str]		    =  one_third - 1
  elif array[two_third] < target:
   __a			:  List[str]		    =  two_third + 1

  else:
   __a			:  Dict		    =  one_third + 1
   __a			:  Any		    =  two_third - 1
 else:
  return -1




def        __UpperCamelCase					(    lowerCAmelCase__							:    int  ,       lowerCAmelCase__							:    int  ,       lowerCAmelCase__							:    list[int]  ,       lowerCAmelCase__							:    int  ):
 if left < right:
  if right - left < precision:
   return lin_search(lowerCAmelCase__  ,       lowerCAmelCase__  ,       lowerCAmelCase__  ,       lowerCAmelCase__  )
  __a			:  Union[str, Any]		    =  (left + right) // 3 + 1
  __a			:  Optional[Any]		    =  2 * (left + right) // 3 + 1

  if array[one_third] == target:
   return one_third
  elif array[two_third] == target:
   return two_third

  elif target < array[one_third]:
   return rec_ternary_search(lowerCAmelCase__  ,       one_third - 1  ,       lowerCAmelCase__  ,       lowerCAmelCase__  )
  elif array[two_third] < target:
   return rec_ternary_search(two_third + 1  ,       lowerCAmelCase__  ,       lowerCAmelCase__  ,       lowerCAmelCase__  )
  else:
   return rec_ternary_search(one_third + 1  ,       two_third - 1  ,       lowerCAmelCase__  ,       lowerCAmelCase__  )
 else:
  return -1


if __name__ == "__main__":
      import doctest

      doctest.testmod()

      lowercase__          =input('Enter numbers separated by comma:\n').strip()
      lowercase__          =[int(item.strip()) for item in user_input.split(',')]
      assert collection == sorted(collection), F"List must be ordered.\n{collection}."
      lowercase__          =int(input('Enter the number to be found in the list:\n').strip())
      lowercase__          =ite_ternary_search(collection, target)
      lowercase__          =rec_ternary_search(0, len(collection) - 1, collection, target)
      if resulta != -1:
            print(F"""Iterative search: {target} found at positions: {resulta}""")
            print(F"""Recursive search: {target} found at positions: {resulta}""")
      else:
            print('Not found')
