import argparse

import gdown
import numpy as np
import torch
from huggingface_hub import hf_hub_download

from transformers import (
    CLIPTokenizer,
    CLIPTokenizerFast,
    VideoMAEImageProcessor,
    XCLIPConfig,
    XCLIPModel,
    XCLIPProcessor,
    XCLIPTextConfig,
    XCLIPVisionConfig,
)


def 	lowerCamelCase					(	UpperCamelCase     :				Any			,       UpperCamelCase     :				Dict							)   ->    Any:
		_lowerCamelCase				   =  XCLIPTextConfig()

		# derive patch size from model name
		_lowerCamelCase				   =  model_name.find('patch'							)
		_lowerCamelCase				   =  int(model_name[start_idx + len('patch'							) : start_idx + len('patch'							) + 2]							)
		_lowerCamelCase				   =  XCLIPVisionConfig(patch_size=UpperCamelCase			,       num_frames=UpperCamelCase							)

		if "large" in model_name:
				_lowerCamelCase				   =  7_68
				_lowerCamelCase				   =  30_72
				_lowerCamelCase				   =  12

				_lowerCamelCase				   =  10_24
				_lowerCamelCase				   =  40_96
				_lowerCamelCase				   =  16
				_lowerCamelCase				   =  24
				_lowerCamelCase				   =  7_68
				_lowerCamelCase				   =  30_72

		if model_name == "xclip-large-patch14-16-frames":
				_lowerCamelCase				   =  3_36

		_lowerCamelCase				   =  XCLIPConfig.from_text_vision_configs(UpperCamelCase			,       UpperCamelCase							)

		if "large" in model_name:
				_lowerCamelCase				   =  7_68

		return config


def 	lowerCamelCase					(	UpperCamelCase     :				Optional[int]							)   ->    int:
		# text encoder
		if name == "token_embedding.weight":
				_lowerCamelCase				   =  name.replace('token_embedding.weight'			,       'text_model.embeddings.token_embedding.weight'							)
		if name == "positional_embedding":
				_lowerCamelCase				   =  name.replace('positional_embedding'			,       'text_model.embeddings.position_embedding.weight'							)
		if "ln_1" in name:
				_lowerCamelCase				   =  name.replace('ln_1'			,       'layer_norm1'							)
		if "ln_2" in name:
				_lowerCamelCase				   =  name.replace('ln_2'			,       'layer_norm2'							)
		if "c_fc" in name:
				_lowerCamelCase				   =  name.replace('c_fc'			,       'fc1'							)
		if "c_proj" in name:
				_lowerCamelCase				   =  name.replace('c_proj'			,       'fc2'							)
		if name.startswith('transformer.resblocks'							):
				_lowerCamelCase				   =  name.replace('transformer.resblocks'			,       'text_model.encoder.layers'							)
		if "attn.out_proj" in name and "message" not in name:
				_lowerCamelCase				   =  name.replace('attn.out_proj'			,       'self_attn.out_proj'							)
		if "ln_final" in name:
				_lowerCamelCase				   =  name.replace('ln_final'			,       'text_model.final_layer_norm'							)
		# visual encoder
		if name == "visual.class_embedding":
				_lowerCamelCase				   =  name.replace('visual.class_embedding'			,       'vision_model.embeddings.class_embedding'							)
		if name == "visual.positional_embedding":
				_lowerCamelCase				   =  name.replace('visual.positional_embedding'			,       'vision_model.embeddings.position_embedding.weight'							)
		if name.startswith('visual.transformer.resblocks'							):
				_lowerCamelCase				   =  name.replace('visual.transformer.resblocks'			,       'vision_model.encoder.layers'							)
		if "visual.conv1" in name:
				_lowerCamelCase				   =  name.replace('visual.conv1'			,       'vision_model.embeddings.patch_embedding'							)
		if "visual.ln_pre" in name:
				_lowerCamelCase				   =  name.replace('visual.ln_pre'			,       'vision_model.pre_layernorm'							)
		if "visual.ln_post" in name:
				_lowerCamelCase				   =  name.replace('visual.ln_post'			,       'vision_model.post_layernorm'							)
		if "visual.proj" in name:
				_lowerCamelCase				   =  name.replace('visual.proj'			,       'visual_projection.weight'							)
		if "text_projection" in name:
				_lowerCamelCase				   =  name.replace('text_projection'			,       'text_projection.weight'							)
		# things on top
		if "prompts_visual_proj" in name:
				_lowerCamelCase				   =  name.replace('prompts_visual_proj'			,       'prompts_visual_projection'							)
		if "prompts_visual_ln" in name:
				_lowerCamelCase				   =  name.replace('prompts_visual_ln'			,       'prompts_visual_layernorm'							)
		# mit
		if name == "mit.positional_embedding":
				_lowerCamelCase				   =  name.replace('positional'			,       'position'							)
		if name.startswith('mit.resblocks'							):
				_lowerCamelCase				   =  name.replace('mit.resblocks'			,       'mit.encoder.layers'							)
		# prompts generator
		if name.startswith('prompts_generator.norm'							):
				_lowerCamelCase				   =  name.replace('prompts_generator.norm'			,       'prompts_generator.layernorm'							)

		return name


def 	lowerCamelCase					(	UpperCamelCase     :				Union[str, Any]			,       UpperCamelCase     :				Optional[Any]							)   ->    Any:
		for key in orig_state_dict.copy().keys():
				_lowerCamelCase				   =  orig_state_dict.pop(UpperCamelCase							)

				if "attn.in_proj" in key:
						_lowerCamelCase				   =  key.split('.'							)
						if key.startswith('visual'							):
								_lowerCamelCase				   =  key_split[3]
								_lowerCamelCase				   =  config.vision_config.hidden_size
								if "message_attn" in key:
										if "weight" in key:
												_lowerCamelCase				   =  val[
												    :dim, :
												]
												_lowerCamelCase				   =  val[
												    dim : dim * 2, :
												]
												_lowerCamelCase				   =  val[
												    -dim:, :
												]
										else:
												_lowerCamelCase				   =  val[
												    :dim
												]
												_lowerCamelCase				   =  val[
												    dim : dim * 2
												]
												_lowerCamelCase				   =  val[
												    -dim:
												]
								else:
										if "weight" in key:
												_lowerCamelCase				   =  val[
												    :dim, :
												]
												_lowerCamelCase				   =  val[
												    dim : dim * 2, :
												]
												_lowerCamelCase				   =  val[
												    -dim:, :
												]
										else:
												_lowerCamelCase				   =  val[:dim]
												_lowerCamelCase				   =  val[
												    dim : dim * 2
												]
												_lowerCamelCase				   =  val[-dim:]
						elif key.startswith('mit'							):
								_lowerCamelCase				   =  key_split[2]
								_lowerCamelCase				   =  config.vision_config.mit_hidden_size
								if "weight" in key:
										_lowerCamelCase				   =  val[:dim, :]
										_lowerCamelCase				   =  val[dim : dim * 2, :]
										_lowerCamelCase				   =  val[-dim:, :]
								else:
										_lowerCamelCase				   =  val[:dim]
										_lowerCamelCase				   =  val[dim : dim * 2]
										_lowerCamelCase				   =  val[-dim:]
						else:
								_lowerCamelCase				   =  key_split[2]
								_lowerCamelCase				   =  config.text_config.hidden_size
								if "weight" in key:
										_lowerCamelCase				   =  val[:dim, :]
										_lowerCamelCase				   =  val[
										    dim : dim * 2, :
										]
										_lowerCamelCase				   =  val[-dim:, :]
								else:
										_lowerCamelCase				   =  val[:dim]
										_lowerCamelCase				   =  val[
										    dim : dim * 2
										]
										_lowerCamelCase				   =  val[-dim:]
				else:
						_lowerCamelCase				   =  rename_key(UpperCamelCase							)
						if new_key_name in ["visual_projection.weight", "text_projection.weight"]:
								_lowerCamelCase				   =  val.T
						_lowerCamelCase				   =  val

		return orig_state_dict


def 	lowerCamelCase					(	UpperCamelCase     :				str							)   ->    Union[str, Any]:
		if num_frames == 8:
				_lowerCamelCase				   =  'eating_spaghetti_8_frames.npy'
		elif num_frames == 16:
				_lowerCamelCase				   =  'eating_spaghetti.npy'
		elif num_frames == 32:
				_lowerCamelCase				   =  'eating_spaghetti_32_frames.npy'
		_lowerCamelCase				   =  hf_hub_download(
		    repo_id='hf-internal-testing/spaghetti-video'			,       filename=UpperCamelCase			,       repo_type='dataset'			,       )
		_lowerCamelCase				   =  np.load(UpperCamelCase							)
		return list(UpperCamelCase							)


def 	lowerCamelCase					(	UpperCamelCase     :				int			,       UpperCamelCase     :				int=None			,       UpperCamelCase     :				List[Any]=False							)   ->    Tuple:
		_lowerCamelCase				   =  {
		    # fully supervised kinetics-400 checkpoints
		    'xclip-base-patch32': 'https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k400_32_8.pth',
		    'xclip-base-patch32-16-frames': (
		        'https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k400_32_16.pth'
		    ),
		    'xclip-base-patch16': 'https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k400_16_8.pth',
		    'xclip-base-patch16-16-frames': (
		        'https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k400_16_16.pth'
		    ),
		    'xclip-large-patch14': 'https://drive.google.com/u/0/uc?id=1NUOImq0o5DlQTST17iIP3vG7DgmHQuCx&amp;export=download&amp;confirm=t&amp;uuid=b26caedc-88e2-473e-830a-9d158b653cdb',
		    'xclip-large-patch14-16-frames': 'https://drive.google.com/u/0/uc?id=1FOYgnJc097OJ4lGwtRCCydQyVPJEOH7d&amp;export=download&amp;confirm=t&amp;uuid=538fa810-e671-4050-b385-9a623f89804f',
		    # fully supervised kinetics-600 checkpoints
		    'xclip-base-patch16-kinetics-600': (
		        'https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k600_16_8.pth'
		    ),
		    'xclip-base-patch16-kinetics-600-16-frames': (
		        'https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k600_16_16.pth'
		    ),
		    'xclip-large-patch14-kinetics-600': 'https://drive.google.com/u/0/uc?id=1FV8C1INuM91sLAN4ImjzePLIlpMSihwV&amp;export=download&amp;confirm=t&amp;uuid=141d4977-4a65-44ae-864f-4b0c19f838be',
		    # few shot
		    'xclip-base-patch16-hmdb-2-shot': (
		        'https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_hmdb_2.pth'
		    ),
		    'xclip-base-patch16-hmdb-4-shot': (
		        'https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_hmdb_4.pth'
		    ),
		    'xclip-base-patch16-hmdb-8-shot': (
		        'https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_hmdb_8.pth'
		    ),
		    'xclip-base-patch16-hmdb-16-shot': (
		        'https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_hmdb_16.pth'
		    ),
		    'xclip-base-patch16-ucf-2-shot': (
		        'https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_ucf_2.pth'
		    ),
		    'xclip-base-patch16-ucf-4-shot': (
		        'https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_ucf_4.pth'
		    ),
		    'xclip-base-patch16-ucf-8-shot': (
		        'https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_ucf_8.pth'
		    ),
		    'xclip-base-patch16-ucf-16-shot': (
		        'https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_ucf_16.pth'
		    ),
		    # zero shot
		    'xclip-base-patch16-zero-shot': 'https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/zero.pth',
		}

		_lowerCamelCase				   =  model_to_url[model_name]
		_lowerCamelCase				   =  8
		if "16-frames" in model_name:
				_lowerCamelCase				   =  16
		elif "shot" in model_name:
				_lowerCamelCase				   =  32

		_lowerCamelCase				   =  get_xclip_config(UpperCamelCase			,       UpperCamelCase							)
		_lowerCamelCase				   =  XCLIPModel(UpperCamelCase							)
		model.eval()

		if "drive" in checkpoint_url:
				_lowerCamelCase				   =  'pytorch_model.bin'
				gdown.cached_download(UpperCamelCase			,       UpperCamelCase			,       quiet=UpperCamelCase							)
				_lowerCamelCase				   =  torch.load(UpperCamelCase			,       map_location='cpu'							)['model']
		else:
				_lowerCamelCase				   =  torch.hub.load_state_dict_from_url(UpperCamelCase							)['model']

		_lowerCamelCase				   =  convert_state_dict(UpperCamelCase			,       UpperCamelCase							)

		_lowerCamelCase				   =  XCLIPModel(UpperCamelCase							)
		_lowerCamelCase   ,							_lowerCamelCase				   =  model.load_state_dict(UpperCamelCase			,       strict=UpperCamelCase							)
		assert missing_keys == ["text_model.embeddings.position_ids", "vision_model.embeddings.position_ids"]
		model.eval()

		_lowerCamelCase				   =  3_36 if model_name == 'xclip-large-patch14-16-frames' else 2_24
		_lowerCamelCase				   =  VideoMAEImageProcessor(size=UpperCamelCase							)
		_lowerCamelCase				   =  CLIPTokenizer.from_pretrained('openai/clip-vit-base-patch32'							)
		_lowerCamelCase				   =  CLIPTokenizerFast.from_pretrained('openai/clip-vit-base-patch32'							)
		_lowerCamelCase				   =  XCLIPProcessor(image_processor=UpperCamelCase			,       tokenizer=UpperCamelCase							)

		_lowerCamelCase				   =  prepare_video(UpperCamelCase							)
		_lowerCamelCase				   =  processor(
		    text=['playing sports', 'eating spaghetti', 'go shopping']			,       videos=UpperCamelCase			,       return_tensors='pt'			,       padding=UpperCamelCase							)

		print('Shape of pixel values:'			,       inputs.pixel_values.shape							)

		with torch.no_grad():
				_lowerCamelCase				   =  model(**UpperCamelCase							)

		# Verify outputs
		_lowerCamelCase				   =  outputs.logits_per_video
		_lowerCamelCase				   =  logits_per_video.softmax(dim=1							)
		print('Probs:'			,       UpperCamelCase							)
		# kinetics-400
		if model_name == "xclip-base-patch32":
				_lowerCamelCase				   =  torch.tensor([[0.0_019, 0.9_951, 0.0_030]]							)
		elif model_name == "xclip-base-patch32-16-frames":
				_lowerCamelCase				   =  torch.tensor([[7.0_9_9_9e-0_4, 9.9_8_8_3e-0_1, 4.5_5_8_0e-0_4]]							)
		elif model_name == "xclip-base-patch16":
				_lowerCamelCase				   =  torch.tensor([[0.0_083, 0.9_681, 0.0_236]]							)
		elif model_name == "xclip-base-patch16-16-frames":
				_lowerCamelCase				   =  torch.tensor([[7.6_9_3_7e-0_4, 9.9_7_2_8e-0_1, 1.9_4_7_3e-0_3]]							)
		elif model_name == "xclip-large-patch14":
				_lowerCamelCase				   =  torch.tensor([[0.0_062, 0.9_864, 0.0_075]]							)
		elif model_name == "xclip-large-patch14-16-frames":
				_lowerCamelCase				   =  torch.tensor([[3.3_8_7_7e-0_4, 9.9_9_3_7e-0_1, 2.8_8_8_8e-0_4]]							)
		# kinetics-600
		elif model_name == "xclip-base-patch16-kinetics-600":
				_lowerCamelCase				   =  torch.tensor([[0.0_555, 0.8_914, 0.0_531]]							)
		elif model_name == "xclip-base-patch16-kinetics-600-16-frames":
				_lowerCamelCase				   =  torch.tensor([[3.8_5_5_4e-0_4, 9.9_9_2_9e-0_1, 3.2_7_5_4e-0_4]]							)
		elif model_name == "xclip-large-patch14-kinetics-600":
				_lowerCamelCase				   =  torch.tensor([[0.0_036, 0.9_920, 0.0_045]]							)
		# few shot
		elif model_name == "xclip-base-patch16-hmdb-2-shot":
				_lowerCamelCase				   =  torch.tensor([[7.1_8_9_0e-0_6, 9.9_9_9_4e-0_1, 5.6_5_5_9e-0_5]]							)
		elif model_name == "xclip-base-patch16-hmdb-4-shot":
				_lowerCamelCase				   =  torch.tensor([[1.0_3_2_0e-0_5, 9.9_9_9_3e-0_1, 6.2_4_3_5e-0_5]]							)
		elif model_name == "xclip-base-patch16-hmdb-8-shot":
				_lowerCamelCase				   =  torch.tensor([[4.1_3_7_7e-0_6, 9.9_9_9_0e-0_1, 9.8_3_8_6e-0_5]]							)
		elif model_name == "xclip-base-patch16-hmdb-16-shot":
				_lowerCamelCase				   =  torch.tensor([[4.1_3_4_7e-0_5, 9.9_9_6_2e-0_1, 3.3_4_1_1e-0_4]]							)
		elif model_name == "xclip-base-patch16-ucf-2-shot":
				_lowerCamelCase				   =  torch.tensor([[8.5_8_5_7e-0_5, 9.9_9_2_8e-0_1, 6.3_2_9_1e-0_4]]							)
		elif model_name == "xclip-base-patch16-ucf-4-shot":
				_lowerCamelCase				   =  torch.tensor([[8.5_8_5_7e-0_5, 9.9_9_2_8e-0_1, 6.3_2_9_1e-0_4]]							)
		elif model_name == "xclip-base-patch16-ucf-8-shot":
				_lowerCamelCase				   =  torch.tensor([[0.0_027, 0.9_904, 0.0_070]]							)
		elif model_name == "xclip-base-patch16-ucf-16-shot":
				_lowerCamelCase				   =  torch.tensor([[9.8_2_1_9e-0_4, 9.9_5_9_3e-0_1, 3.0_8_6_3e-0_3]]							)
		# zero shot
		elif model_name == "xclip-base-patch16-zero-shot":
				_lowerCamelCase				   =  torch.tensor([[3.5_0_8_2e-0_4, 9.9_7_8_5e-0_1, 1.7_9_6_6e-0_3]]							)
		else:
				raise ValueError(F"""Model name {model_name} not supported"""							)
		assert torch.allclose(UpperCamelCase			,       UpperCamelCase			,       atol=1e-3							)
		print('Looks ok!'							)

		if pytorch_dump_folder_path is not None:
				print(F"""Saving model {model_name} to {pytorch_dump_folder_path}"""							)
				model.save_pretrained(UpperCamelCase							)

		if push_to_hub:
				print('Pushing model, processor and slow tokenizer files to the hub...'							)
				model.push_to_hub(UpperCamelCase			,       organization='nielsr'							)
				processor.push_to_hub(UpperCamelCase			,       organization='nielsr'							)
				slow_tokenizer.push_to_hub(UpperCamelCase			,       organization='nielsr'							)


if __name__ == "__main__":
				A						=			argparse.ArgumentParser()
				# Required parameters
				parser.add_argument(
				    '--model_name',
				    default='xclip-base-patch32',
				    type=str,
				    help='Name of the model.',
				)
				parser.add_argument(
				    '--pytorch_dump_folder_path', default=None, type=str, help='Path to the output PyTorch model directory.'
				)
				parser.add_argument(
				    '--push_to_hub', action='store_true', help='Whether or not to push the converted model to the 🤗 hub.'
				)

				A						=			parser.parse_args()
				convert_xclip_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)
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     lowerCAmelCase__     (				SCREAMING_SNAKE_CASE_     ,unittest.TestCase			):


		'''simple docstring'''






		lowerCAmelCase_     =   ShapEImgaImgPipeline
		lowerCAmelCase_     =   ['image']
		lowerCAmelCase_     =   ['image']
		lowerCAmelCase_     =   [
		    'num_images_per_prompt',
		    'num_inference_steps',
		    'generator',
		    'latents',
		    'guidance_scale',
		    'frame_size',
		    'output_type',
		    'return_dict',
		]
		lowerCAmelCase_     =   False
		@property
		def   _snake_case       (       self     :     List[str]		)      ->	Any:
				return 3_2
		@property
		def   _snake_case       (       self     :     int		)      ->	Dict:
				return 3_2
		@property
		def   _snake_case       (       self     :     Tuple		)      ->	Optional[Any]:
				return self.time_input_dim * 4
		@property
		def   _snake_case       (       self     :     Tuple		)      ->	Optional[Any]:
				return 8
		@property
		def   _snake_case       (       self     :     int		)      ->	Dict:
				torch.manual_seed(0		)
				_lowerCamelCase				   =  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				   =  CLIPVisionModel(snake_case__		)
				return model
		@property
		def   _snake_case       (       self     :     Tuple		)      ->	Any:
				_lowerCamelCase				   =  CLIPImageProcessor(
				    crop_size=2_2_4     ,  do_center_crop=snake_case__     ,  do_normalize=snake_case__     ,  do_resize=snake_case__     ,  image_mean=[0.4814_5466, 0.457_8275, 0.4082_1073]     ,  image_std=[0.2686_2954, 0.2613_0258, 0.2757_7711]     ,  resample=3     ,  size=2_2_4     ,  )

				return image_processor
		@property
		def   _snake_case       (       self     :     List[str]		)      ->	Any:
				torch.manual_seed(0		)

				_lowerCamelCase				   =  {
				    '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				   =  PriorTransformer(**snake_case__		)
				return model
		@property
		def   _snake_case       (       self     :     List[Any]		)      ->	Optional[int]:
				torch.manual_seed(0		)

				_lowerCamelCase				   =  {
				    '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				   =  ShapERenderer(**snake_case__		)
				return model
		def   _snake_case       (       self     :     Any		)      ->	str:
				_lowerCamelCase				   =  self.dummy_prior
				_lowerCamelCase				   =  self.dummy_image_encoder
				_lowerCamelCase				   =  self.dummy_image_processor
				_lowerCamelCase				   =  self.dummy_renderer

				_lowerCamelCase				   =  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     ,  )
				_lowerCamelCase				   =  {
				    'prior': prior,
				    'image_encoder': image_encoder,
				    'image_processor': image_processor,
				    'renderer': renderer,
				    'scheduler': scheduler,
				}

				return components
		def   _snake_case       (       self     :     List[str]     ,  snake_case__     :     str     ,  snake_case__     :     Optional[Any]=0		)      ->	Union[str, Any]:
				_lowerCamelCase				   =  floats_tensor((1, 3, 6_4, 6_4)     ,  rng=random.Random(snake_case__		)		).to(snake_case__		)

				if str(snake_case__		).startswith('mps'		):
						_lowerCamelCase				   =  torch.manual_seed(snake_case__		)
				else:
						_lowerCamelCase				   =  torch.Generator(device=snake_case__		).manual_seed(snake_case__		)
				_lowerCamelCase				   =  {
				    'image': input_image,
				    'generator': generator,
				    'num_inference_steps': 1,
				    'frame_size': 3_2,
				    'output_type': 'np',
				}
				return inputs
		def   _snake_case       (       self     :     Union[str, Any]		)      ->	Union[str, Any]:
				_lowerCamelCase				   =  'cpu'

				_lowerCamelCase				   =  self.get_dummy_components()

				_lowerCamelCase				   =  self.pipeline_class(**snake_case__		)
				_lowerCamelCase				   =  pipe.to(snake_case__		)

				pipe.set_progress_bar_config(disable=snake_case__		)

				_lowerCamelCase				   =  pipe(**self.get_dummy_inputs(snake_case__		)		)
				_lowerCamelCase				   =  output.images[0]
				_lowerCamelCase				   =  image[0, -3:, -3:, -1]

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

				_lowerCamelCase				   =  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   _snake_case       (       self     :     Dict		)      ->	Optional[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   _snake_case       (       self     :     Optional[Any]		)      ->	Tuple:
				_lowerCamelCase				   =  torch_device == 'cpu'
				_lowerCamelCase				   =  True
				self._test_inference_batch_single_identical(
				    batch_size=2     ,  test_max_difference=snake_case__     ,  relax_max_difference=snake_case__     ,  )



		def   _snake_case       (       self     :     Any		)      ->	Tuple:
				_lowerCamelCase				   =  self.get_dummy_components()
				_lowerCamelCase				   =  self.pipeline_class(**snake_case__		)
				_lowerCamelCase				   =  pipe.to(snake_case__		)
				pipe.set_progress_bar_config(disable=snake_case__		)

				_lowerCamelCase				   =  1
				_lowerCamelCase				   =  2

				_lowerCamelCase				   =  self.get_dummy_inputs(snake_case__		)

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

				_lowerCamelCase				   =  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     lowerCAmelCase__     (				unittest.TestCase			):


		'''simple docstring'''






		def   _snake_case       (       self     :     Union[str, Any]		)      ->	Optional[int]:
				# clean up the VRAM after each test
				super().tearDown()
				gc.collect()
				torch.cuda.empty_cache()



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

				_lowerCamelCase				   =  torch.Generator(device=snake_case__		).manual_seed(0		)

				_lowerCamelCase				   =  pipe(
				    snake_case__     ,  generator=snake_case__     ,  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(snake_case__     ,  snake_case__		)