



'''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_				(       snake_case_     ,					unittest.TestCase      ):
						_lowerCAmelCase  : Any						  =					ShapEPipeline
						_lowerCAmelCase  : List[str]						  =					['prompt']
						_lowerCAmelCase  : List[Any]						  =					['prompt']
						_lowerCAmelCase  : Union[str, Any]						  =					[
						    'num_images_per_prompt',
						    'num_inference_steps',
						    'generator',
						    'latents',
						    'guidance_scale',
						    'frame_size',
						    'output_type',
						    'return_dict',
						]
						_lowerCAmelCase  : Union[str, Any]						  =					False







						@property
						def        __lowercase     (	self :   Optional[Any]					):

													"""simple docstring"""
													return 32







						@property
						def        __lowercase     (	self :   Any					):

													"""simple docstring"""
													return 32







						@property
						def        __lowercase     (	self :   int					):

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







						@property
						def        __lowercase     (	self :   Optional[int]					):

													"""simple docstring"""
													return 8







						@property
						def        __lowercase     (	self :   Tuple					):

													"""simple docstring"""
													SCREAMING_SNAKE_CASE       :   int         =      CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''					)
													return tokenizer







						@property
						def        __lowercase     (	self :   Any					):

													"""simple docstring"""
													torch.manual_seed(0					)
													SCREAMING_SNAKE_CASE       :   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=37					,     layer_norm_eps=1e-05					,     num_attention_heads=4					,     num_hidden_layers=5					,     pad_token_id=1					,     vocab_size=10_00					,     )
													return CLIPTextModelWithProjection(lowerCAmelCase__					)







						@property
						def        __lowercase     (	self :   Any					):

													"""simple docstring"""
													torch.manual_seed(0					)

													SCREAMING_SNAKE_CASE       :   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,
													}

													SCREAMING_SNAKE_CASE       :   str         =      PriorTransformer(**lowerCAmelCase__					)
													return model







						@property
						def        __lowercase     (	self :   Any					):

													"""simple docstring"""
													torch.manual_seed(0					)

													SCREAMING_SNAKE_CASE       :   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       :   Tuple         =      ShapERenderer(**lowerCAmelCase__					)
													return model







						def        __lowercase     (	self :   int					):

													"""simple docstring"""
													SCREAMING_SNAKE_CASE       :   Tuple         =      self.dummy_prior
													SCREAMING_SNAKE_CASE       :   List[Any]         =      self.dummy_text_encoder
													SCREAMING_SNAKE_CASE       :   List[Any]         =      self.dummy_tokenizer
													SCREAMING_SNAKE_CASE       :   List[Any]         =      self.dummy_renderer

													SCREAMING_SNAKE_CASE       :   List[str]         =      HeunDiscreteScheduler(
													    beta_schedule='''exp'''					,     num_train_timesteps=10_24					,     prediction_type='''sample'''					,     use_karras_sigmas=lowerCAmelCase__					,     clip_sample=lowerCAmelCase__					,     clip_sample_range=1.0					,     )
													SCREAMING_SNAKE_CASE       :   Tuple         =      {
													    '''prior''': prior,
													    '''text_encoder''': text_encoder,
													    '''tokenizer''': tokenizer,
													    '''renderer''': renderer,
													    '''scheduler''': scheduler,
													}

													return components







						def        __lowercase     (	self :   str					,     lowerCAmelCase__ :   Tuple					,     lowerCAmelCase__ :   int=0					):

													"""simple docstring"""
													if str(lowerCAmelCase__					).startswith('''mps'''					):
																				SCREAMING_SNAKE_CASE       :   Any         =      torch.manual_seed(lowerCAmelCase__					)
													else:
																				SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      torch.Generator(device=lowerCAmelCase__					).manual_seed(lowerCAmelCase__					)
													SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      {
													    '''prompt''': '''horse''',
													    '''generator''': generator,
													    '''num_inference_steps''': 1,
													    '''frame_size''': 32,
													    '''output_type''': '''np''',
													}
													return inputs







						def        __lowercase     (	self :   str					):

													"""simple docstring"""
													SCREAMING_SNAKE_CASE       :   Dict         =      '''cpu'''

													SCREAMING_SNAKE_CASE       :   int         =      self.get_dummy_components()

													SCREAMING_SNAKE_CASE       :   List[str]         =      self.pipeline_class(**lowerCAmelCase__					)
													SCREAMING_SNAKE_CASE       :   int         =      pipe.to(lowerCAmelCase__					)

													pipe.set_progress_bar_config(disable=lowerCAmelCase__					)

													SCREAMING_SNAKE_CASE       :   List[Any]         =      pipe(**self.get_dummy_inputs(lowerCAmelCase__					)					)
													SCREAMING_SNAKE_CASE       :   Optional[Any]         =      output.images[0]
													SCREAMING_SNAKE_CASE       :   List[str]         =      image[0, -3:, -3:, -1]

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

													SCREAMING_SNAKE_CASE       :   int         =      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        __lowercase     (	self :   Tuple					):

													"""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        __lowercase     (	self :   Union[str, Any]					):

													"""simple docstring"""
													SCREAMING_SNAKE_CASE       :   int         =      torch_device == '''cpu'''
													SCREAMING_SNAKE_CASE       :   int         =      True

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







						def        __lowercase     (	self :   Dict					):

													"""simple docstring"""
													SCREAMING_SNAKE_CASE       :   List[Any]         =      self.get_dummy_components()
													SCREAMING_SNAKE_CASE       :   Any         =      self.pipeline_class(**lowerCAmelCase__					)
													SCREAMING_SNAKE_CASE       :   Optional[Any]         =      pipe.to(lowerCAmelCase__					)
													pipe.set_progress_bar_config(disable=lowerCAmelCase__					)

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

													SCREAMING_SNAKE_CASE       :   Optional[int]         =      self.get_dummy_inputs(lowerCAmelCase__					)

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

													SCREAMING_SNAKE_CASE       :   str         =      pipe(**lowerCAmelCase__					,     num_images_per_prompt=lowerCAmelCase__					)[0]

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


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







						def        __lowercase     (	self :   Optional[int]					):

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







						def        __lowercase     (	self :   Tuple					):

													"""simple docstring"""
													SCREAMING_SNAKE_CASE       :   List[str]         =      load_numpy(
													    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
													    '''/shap_e/test_shap_e_np_out.npy'''					)
													SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      ShapEPipeline.from_pretrained('''openai/shap-e'''					)
													SCREAMING_SNAKE_CASE       :   int         =      pipe.to(lowerCAmelCase__					)
													pipe.set_progress_bar_config(disable=lowerCAmelCase__					)

													SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      torch.Generator(device=lowerCAmelCase__					).manual_seed(0					)

													SCREAMING_SNAKE_CASE       :   Tuple         =      pipe(
													    '''a shark'''					,     generator=lowerCAmelCase__					,     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(lowerCAmelCase__					,     lowerCAmelCase__					)









'''simple docstring'''



import unittest

from transformers import is_tf_available
from transformers.testing_utils import require_tf


if is_tf_available():
							import tensorflow as tf
							from tensorflow.python.eager import context
							from tensorflow.python.framework import ops

							from transformers import GradientAccumulator, create_optimizer
@require_tf
class 						lowerCamelCase_				(       unittest.TestCase      ):







						def        __lowercase     (	self :   List[Any]					,     lowerCAmelCase__ :   str					,     lowerCAmelCase__ :   Dict					,     lowerCAmelCase__ :   int					):

													"""simple docstring"""
													self.assertEqual(len(lowerCAmelCase__					)					,     len(lowerCAmelCase__					)					)
													for a, b in zip(lowerCAmelCase__					,     lowerCAmelCase__					):
																				self.assertAlmostEqual(lowerCAmelCase__					,     lowerCAmelCase__					,     delta=lowerCAmelCase__					)







						def        __lowercase     (	self :   Dict					):

													"""simple docstring"""
													SCREAMING_SNAKE_CASE       :   List[Any]         =      GradientAccumulator()
													accumulator([tf.constant([1.0, 2.0]					)]					)
													accumulator([tf.constant([-2.0, 1.0]					)]					)
													accumulator([tf.constant([-1.0, 2.0]					)]					)
													with self.assertRaises(lowerCAmelCase__					):
																				accumulator([tf.constant([1.0, 1.0]					), tf.constant([2.0, 2.0]					)]					)
													self.assertEqual(accumulator.step					,     3					)
													self.assertEqual(len(accumulator.gradients					)					,     1					)
													self.assertListAlmostEqual(accumulator.gradients[0].numpy().tolist()					,     [-2.0, 5.0]					,     tol=1e-2					)
													accumulator.reset()
													self.assertEqual(accumulator.step					,     0					)
													self.assertListAlmostEqual(accumulator.gradients[0].numpy().tolist()					,     [0.0, 0.0]					,     tol=1e-2					)







						def        __lowercase     (	self :   List[str]					):

													"""simple docstring"""
													SCREAMING_SNAKE_CASE       :   Optional[int]         =      None
													ops.enable_eager_execution_internal()
													SCREAMING_SNAKE_CASE       :   List[Any]         =      tf.config.list_physical_devices('''CPU'''					)
													if len(lowerCAmelCase__					) == 1:
																				tf.config.set_logical_device_configuration(
																				    physical_devices[0]					,     [tf.config.LogicalDeviceConfiguration(), tf.config.LogicalDeviceConfiguration()]					)
													SCREAMING_SNAKE_CASE       :   Tuple         =      tf.config.list_logical_devices(device_type='''CPU'''					)
													SCREAMING_SNAKE_CASE       :   List[Any]         =      tf.distribute.MirroredStrategy(devices=devices[:2]					)

													with strategy.scope():
																				SCREAMING_SNAKE_CASE       :   Optional[Any]         =      GradientAccumulator()
																				SCREAMING_SNAKE_CASE       :   Dict         =      tf.Variable([4.0, 3.0]					)
																				SCREAMING_SNAKE_CASE       ,SCREAMING_SNAKE_CASE       :   Optional[int]         =      create_optimizer(5e-5					,     10					,     5					)
																				SCREAMING_SNAKE_CASE       :   List[Any]         =      tf.Variable([0.0, 0.0]					,     trainable=lowerCAmelCase__					)

													def accumulate_on_replica(lowerCAmelCase__ :   List[str]					):
																				accumulator([gradient]					)

													def apply_on_replica():
																				optimizer.apply_gradients(list(zip(accumulator.gradients					,     [variable]					)					)					)

													@tf.function
													def accumulate(lowerCAmelCase__ :   str					,     lowerCAmelCase__ :   List[Any]					):
																				with strategy.scope():
																											SCREAMING_SNAKE_CASE       :   Optional[Any]         =      strategy.experimental_local_results(lowerCAmelCase__					)
																											local_variables[0].assign(lowerCAmelCase__					)
																											local_variables[1].assign(lowerCAmelCase__					)
																											strategy.run(lowerCAmelCase__					,     args=(gradient_placeholder,)					)

													@tf.function
													def apply_grad():
																				with strategy.scope():
																											strategy.run(lowerCAmelCase__					)

													def _check_local_values(lowerCAmelCase__ :   Dict					,     lowerCAmelCase__ :   int					):
																				SCREAMING_SNAKE_CASE       :   Tuple         =      strategy.experimental_local_results(accumulator._gradients[0]					)
																				self.assertListAlmostEqual(values[0].value()					,     lowerCAmelCase__					,     tol=1e-2					)
																				self.assertListAlmostEqual(values[1].value()					,     lowerCAmelCase__					,     tol=1e-2					)

													accumulate([1.0, 2.0]					,     [-1.0, 1.0]					)
													accumulate([3.0, -1.0]					,     [-1.0, -1.0]					)
													accumulate([-2.0, 2.0]					,     [3.0, -2.0]					)
													self.assertEqual(accumulator.step					,     3					)
													_check_local_values([2.0, 3.0]					,     [1.0, -2.0]					)
													apply_grad()
													self.assertListAlmostEqual(variable.value()					,     [4.0, 3.0]					,     tol=1e-2					)
													accumulator.reset()
													self.assertEqual(accumulator.step					,     0					)
													_check_local_values([0.0, 0.0]					,     [0.0, 0.0]					)




