


import gc
import random
import unittest

import numpy as np
import torch
from transformers import CLIPImageProcessor, CLIPVisionConfig, CLIPVisionModel

from diffusers import HeunDiscreteScheduler, PriorTransformer, ShapEImgaImgPipeline
from diffusers.pipelines.shap_e import ShapERenderer
from diffusers.utils import floats_tensor, load_image, 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__(							UpperCAmelCase__,						unittest.TestCase							):







						'''simple docstring'''




						SCREAMING_SNAKE_CASE__      :  Optional[Any]					=							ShapEImgaImgPipeline
						SCREAMING_SNAKE_CASE__      :  List[str]					=							["""image"""]
						SCREAMING_SNAKE_CASE__      :  List[str]					=							["""image"""]
						SCREAMING_SNAKE_CASE__      :  int					=							[
						    """num_images_per_prompt""",
						    """num_inference_steps""",
						    """generator""",
						    """latents""",
						    """guidance_scale""",
						    """frame_size""",
						    """output_type""",
						    """return_dict""",
						]
						SCREAMING_SNAKE_CASE__      :  str					=							False




						@property
						def 						lowercase_ (				self							)					->    Any:
										return 3_2




						@property
						def 						lowercase_ (				self							)					->    List[str]:
										return 3_2




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




						@property
						def 						lowercase_ (				self							)					->    Dict:
										return 8




						@property
						def 						lowercase_ (				self							)					->    int:
										torch.manual_seed(0							)
										lowerCAmelCase_ :      List[str]    				=				CLIPVisionConfig(
										    hidden_size=self.text_embedder_hidden_size   ,						image_size=6_4   ,						projection_dim=self.text_embedder_hidden_size   ,						intermediate_size=3_7   ,						num_attention_heads=4   ,						num_channels=3   ,						num_hidden_layers=5   ,						patch_size=1   ,						)

										lowerCAmelCase_ :      int    				=				CLIPVisionModel(__lowercase							)
										return model




						@property
						def 						lowercase_ (				self							)					->    Dict:
										lowerCAmelCase_ :      Union[str, Any]    				=				CLIPImageProcessor(
										    crop_size=2_2_4   ,						do_center_crop=__lowercase   ,						do_normalize=__lowercase   ,						do_resize=__lowercase   ,						image_mean=[0.48_14_54_66, 0.4_57_82_75, 0.40_82_10_73]   ,						image_std=[0.26_86_29_54, 0.26_13_02_58, 0.27_57_77_11]   ,						resample=3   ,						size=2_2_4   ,						)

										return image_processor




						@property
						def 						lowercase_ (				self							)					->    List[str]:
										torch.manual_seed(0							)

										lowerCAmelCase_ :      Optional[int]    				=				{
										    '''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''',
										    '''embedding_proj_norm_type''': '''layer''',
										    '''encoder_hid_proj_type''': None,
										    '''added_emb_type''': None,
										}

										lowerCAmelCase_ :      int    				=				PriorTransformer(**__lowercase							)
										return model




						@property
						def 						lowercase_ (				self							)					->    Dict:
										torch.manual_seed(0							)

										lowerCAmelCase_ :      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,
										    ),
										}
										lowerCAmelCase_ :      Tuple    				=				ShapERenderer(**__lowercase							)
										return model




						def 						lowercase_ (				self							)					->    Any:
										lowerCAmelCase_ :      Optional[Any]    				=				self.dummy_prior
										lowerCAmelCase_ :      int    				=				self.dummy_image_encoder
										lowerCAmelCase_ :      Optional[Any]    				=				self.dummy_image_processor
										lowerCAmelCase_ :      List[str]    				=				self.dummy_renderer

										lowerCAmelCase_ :      Optional[Any]    				=				HeunDiscreteScheduler(
										    beta_schedule='''exp'''   ,						num_train_timesteps=1_0_2_4   ,						prediction_type='''sample'''   ,						use_karras_sigmas=__lowercase   ,						clip_sample=__lowercase   ,						clip_sample_range=1.0   ,						)
										lowerCAmelCase_ :      Tuple    				=				{
										    '''prior''': prior,
										    '''image_encoder''': image_encoder,
										    '''image_processor''': image_processor,
										    '''renderer''': renderer,
										    '''scheduler''': scheduler,
										}

										return components




						def 						lowercase_ (				self   ,						__lowercase   ,						__lowercase=0							)					->    Optional[Any]:
										lowerCAmelCase_ :      int    				=				floats_tensor((1, 3, 6_4, 6_4)   ,						rng=random.Random(__lowercase							)							).to(__lowercase							)

										if str(__lowercase							).startswith('''mps'''							):
														lowerCAmelCase_ :      Union[str, Any]    				=				torch.manual_seed(__lowercase							)
										else:
														lowerCAmelCase_ :      Any    				=				torch.Generator(device=__lowercase							).manual_seed(__lowercase							)
										lowerCAmelCase_ :      Dict    				=				{
										    '''image''': input_image,
										    '''generator''': generator,
										    '''num_inference_steps''': 1,
										    '''frame_size''': 3_2,
										    '''output_type''': '''np''',
										}
										return inputs




						def 						lowercase_ (				self							)					->    int:
										lowerCAmelCase_ :      Optional[int]    				=				'''cpu'''

										lowerCAmelCase_ :      Tuple    				=				self.get_dummy_components()

										lowerCAmelCase_ :      Any    				=				self.pipeline_class(**__lowercase							)
										lowerCAmelCase_ :      Optional[int]    				=				pipe.to(__lowercase							)

										pipe.set_progress_bar_config(disable=__lowercase							)

										lowerCAmelCase_ :      Union[str, Any]    				=				pipe(**self.get_dummy_inputs(__lowercase							)							)
										lowerCAmelCase_ :      Optional[int]    				=				output.images[0]
										lowerCAmelCase_ :      List[str]    				=				image[0, -3:, -3:, -1]

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

										lowerCAmelCase_ :      str    				=				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 						lowercase_ (				self							)					->    Optional[int]:
										# 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							)					->    List[Any]:
										lowerCAmelCase_ :      Any    				=				torch_device == '''cpu'''
										lowerCAmelCase_ :      str    				=				True
										self._test_inference_batch_single_identical(
										    batch_size=2   ,						test_max_difference=__lowercase   ,						relax_max_difference=__lowercase   ,						)




						def 						lowercase_ (				self							)					->    Optional[int]:
										lowerCAmelCase_ :      Union[str, Any]    				=				self.get_dummy_components()
										lowerCAmelCase_ :      Any    				=				self.pipeline_class(**__lowercase							)
										lowerCAmelCase_ :      Optional[Any]    				=				pipe.to(__lowercase							)
										pipe.set_progress_bar_config(disable=__lowercase							)

										lowerCAmelCase_ :      str    				=				1
										lowerCAmelCase_ :      int    				=				2

										lowerCAmelCase_ :      Tuple    				=				self.get_dummy_inputs(__lowercase							)

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

										lowerCAmelCase_ :      List[Any]    				=				pipe(**__lowercase   ,						num_images_per_prompt=__lowercase							)[0]

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







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







						'''simple docstring'''




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




						def 						lowercase_ (				self							)					->    Tuple:
										lowerCAmelCase_ :      List[Any]    				=				load_image(
										    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main''' '''/shap_e/corgi.png'''							)
										lowerCAmelCase_ :      Dict    				=				load_numpy(
										    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
										    '''/shap_e/test_shap_e_img2img_out.npy'''							)
										lowerCAmelCase_ :      Union[str, Any]    				=				ShapEImgaImgPipeline.from_pretrained('''openai/shap-e-img2img'''							)
										lowerCAmelCase_ :      Tuple    				=				pipe.to(__lowercase							)
										pipe.set_progress_bar_config(disable=__lowercase							)

										lowerCAmelCase_ :      str    				=				torch.Generator(device=__lowercase							).manual_seed(0							)

										lowerCAmelCase_ :      Union[str, Any]    				=				pipe(
										    __lowercase   ,						generator=__lowercase   ,						guidance_scale=3.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(__lowercase   ,						__lowercase							)








# We ignore warnings about stepping the scheduler since we step it ourselves during gradient accumulation

import warnings

from .state import AcceleratorState, GradientState


warnings.filterwarnings("""ignore""", category=UserWarning, module="""torch.optim.lr_scheduler""")







class 				snake_case__:







						'''simple docstring'''




						def __init__(				self   ,						__lowercase   ,						__lowercase   ,						__lowercase = True   ,						__lowercase = False							)					->    Tuple:
										lowerCAmelCase_ :      Optional[int]    				=				scheduler
										lowerCAmelCase_ :      Dict    				=				optimizers if isinstance(__lowercase   ,						(list, tuple)							) else [optimizers]
										lowerCAmelCase_ :      str    				=				split_batches
										lowerCAmelCase_ :      Any    				=				step_with_optimizer
										lowerCAmelCase_ :      Optional[Any]    				=				GradientState()




						def 						lowercase_ (				self   ,						*__lowercase   ,						**__lowercase							)					->    Any:
										if not self.step_with_optimizer:
														# No link between scheduler and optimizer -> just step
														self.scheduler.step(*__lowercase   ,						**__lowercase							)
														return

										# Otherwise, first make sure the optimizer was stepped.
										if not self.gradient_state.sync_gradients:
														if self.gradient_state.adjust_scheduler:
																		self.scheduler._step_count += 1
														return

										for opt in self.optimizers:
														if opt.step_was_skipped:
																		return
										if self.split_batches:
														# Split batches -> the training dataloader batch size is not changed so one step per training step
														self.scheduler.step(*__lowercase   ,						**__lowercase							)
										else:
														# Otherwise the training dataloader batch size was multiplied by `num_processes`, so we need to do
														# num_processes steps per training step
														lowerCAmelCase_ :      Optional[Any]    				=				AcceleratorState().num_processes
														for _ in range(__lowercase							):
																		# Special case when using OneCycle and `drop_last` was not used
																		if hasattr(self.scheduler   ,						'''total_steps'''							):
																						if self.scheduler._step_count <= self.scheduler.total_steps:
																										self.scheduler.step(*__lowercase   ,						**__lowercase							)
																		else:
																						self.scheduler.step(*__lowercase   ,						**__lowercase							)




						def 						lowercase_ (				self							)					->    Union[str, Any]:
										return self.scheduler.get_last_lr()




						def 						lowercase_ (				self							)					->    List[str]:
										return self.scheduler.state_dict()




						def 						lowercase_ (				self   ,						__lowercase							)					->    int:
										self.scheduler.load_state_dict(__lowercase							)




						def 						lowercase_ (				self							)					->    Tuple:
										return self.scheduler.get_lr()




						def 						lowercase_ (				self   ,						*__lowercase   ,						**__lowercase							)					->    int:
										return self.scheduler.print_lr(*__lowercase   ,						**__lowercase							)