
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        _SCREAMING_SNAKE_CASE	( snake_case_     ,			unittest.TestCase		):
	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        SCREAMING_SNAKE_CASE_(       self					)							->			Dict:
					return 32



	@property
	def        SCREAMING_SNAKE_CASE_(       self					)							->			Optional[int]:
					return 32



	@property
	def        SCREAMING_SNAKE_CASE_(       self					)							->			Dict:
					return self.time_input_dim * 4



	@property
	def        SCREAMING_SNAKE_CASE_(       self					)							->			str:
					return 8



	@property
	def        SCREAMING_SNAKE_CASE_(       self					)							->			Optional[Any]:
					torch.manual_seed(0					)
					lowerCamelCase_         =     CLIPVisionConfig(
					    hidden_size=self.text_embedder_hidden_size		,    image_size=64		,    projection_dim=self.text_embedder_hidden_size		,    intermediate_size=37		,    num_attention_heads=4		,    num_channels=3		,    num_hidden_layers=5		,    patch_size=1		,    )

					lowerCamelCase_         =     CLIPVisionModel(lowercase					)
					return model



	@property
	def        SCREAMING_SNAKE_CASE_(       self					)							->			List[Any]:
					lowerCamelCase_         =     CLIPImageProcessor(
					    crop_size=224		,    do_center_crop=lowercase		,    do_normalize=lowercase		,    do_resize=lowercase		,    image_mean=[0.4_8_1_4_5_4_6_6, 0.4_5_7_8_2_7_5, 0.4_0_8_2_1_0_7_3]		,    image_std=[0.2_6_8_6_2_9_5_4, 0.2_6_1_3_0_2_5_8, 0.2_7_5_7_7_7_1_1]		,    resample=3		,    size=224		,    )

					return image_processor



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

					lowerCamelCase_         =     {
					    "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",
					    "embedding_proj_norm_type": "layer",
					    "encoder_hid_proj_type": None,
					    "added_emb_type": None,
					}

					lowerCamelCase_         =     PriorTransformer(**lowercase					)
					return model



	@property
	def        SCREAMING_SNAKE_CASE_(       self					)							->			Union[str, Any]:
					torch.manual_seed(0					)

					lowerCamelCase_         =     {
					    "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,
					    ),
					}
					lowerCamelCase_         =     ShapERenderer(**lowercase					)
					return model



	def        SCREAMING_SNAKE_CASE_(       self					)							->			Dict:
					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=1024		,    prediction_type="sample"		,    use_karras_sigmas=lowercase		,    clip_sample=lowercase		,    clip_sample_range=1.0		,    )
					lowerCamelCase_         =     {
					    "prior": prior,
					    "image_encoder": image_encoder,
					    "image_processor": image_processor,
					    "renderer": renderer,
					    "scheduler": scheduler,
					}

					return components



	def        SCREAMING_SNAKE_CASE_(       self		,    lowercase		,    lowercase=0					)							->			int:
					lowerCamelCase_         =     floats_tensor((1, 3, 64, 64)		,    rng=random.Random(lowercase					)					).to(lowercase					)

					if str(lowercase					).startswith("mps"					):
									lowerCamelCase_         =     torch.manual_seed(lowercase					)
					else:
									lowerCamelCase_         =     torch.Generator(device=lowercase					).manual_seed(lowercase					)
					lowerCamelCase_         =     {
					    "image": input_image,
					    "generator": generator,
					    "num_inference_steps": 1,
					    "frame_size": 32,
					    "output_type": "np",
					}
					return inputs



	def        SCREAMING_SNAKE_CASE_(       self					)							->			Tuple:
					lowerCamelCase_         =     "cpu"

					lowerCamelCase_         =     self.get_dummy_components()

					lowerCamelCase_         =     self.pipeline_class(**lowercase					)
					lowerCamelCase_         =     pipe.to(lowercase					)

					pipe.set_progress_bar_config(disable=lowercase					)

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

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

					lowerCamelCase_         =     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        SCREAMING_SNAKE_CASE_(       self					)							->			List[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        SCREAMING_SNAKE_CASE_(       self					)							->			List[Any]:
					lowerCamelCase_         =     torch_device == "cpu"
					lowerCamelCase_         =     True
					self._test_inference_batch_single_identical(
					    batch_size=2		,    test_max_difference=lowercase		,    relax_max_difference=lowercase		,    )



	def        SCREAMING_SNAKE_CASE_(       self					)							->			Tuple:
					lowerCamelCase_         =     self.get_dummy_components()
					lowerCamelCase_         =     self.pipeline_class(**lowercase					)
					lowerCamelCase_         =     pipe.to(lowercase					)
					pipe.set_progress_bar_config(disable=lowercase					)

					lowerCamelCase_         =     1
					lowerCamelCase_         =     2

					lowerCamelCase_         =     self.get_dummy_inputs(lowercase					)

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

					lowerCamelCase_         =     pipe(**lowercase		,    num_images_per_prompt=lowercase					)[0]

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




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



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



	def        SCREAMING_SNAKE_CASE_(       self					)							->			int:
					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(lowercase					)
					pipe.set_progress_bar_config(disable=lowercase					)

					lowerCamelCase_         =     torch.Generator(device=lowercase					).manual_seed(0					)

					lowerCamelCase_         =     pipe(
					    lowercase		,    generator=lowercase		,    guidance_scale=3.0		,    num_inference_steps=64		,    frame_size=64		,    output_type="np"		,    ).images[0]

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

					assert_mean_pixel_difference(lowercase		,    lowercase					)




import warnings
from contextlib import contextmanager

from ....processing_utils import ProcessorMixin




class        _SCREAMING_SNAKE_CASE	( snake_case_		):
	lowerCAmelCase__   =					'MCTCTFeatureExtractor'
	lowerCAmelCase__   =					'AutoTokenizer'



	def __init__(       self		,    lowercase		,    lowercase					)							->			str:
					super().__init__(lowercase		,    lowercase					)
					lowerCamelCase_         =     self.feature_extractor
					lowerCamelCase_         =     False



	def __call__(       self		,    *lowercase		,    **lowercase					)							->			List[str]:
					# For backward compatibility
					if self._in_target_context_manager:
									return self.current_processor(*lowercase		,    **lowercase					)

					if "raw_speech" in kwargs:
									warnings.warn("Using `raw_speech` as a keyword argument is deprecated. Use `audio` instead."					)
									lowerCamelCase_         =     kwargs.pop("raw_speech"					)
					else:
									lowerCamelCase_         =     kwargs.pop("audio"		,    lowercase					)
					lowerCamelCase_         =     kwargs.pop("sampling_rate"		,    lowercase					)
					lowerCamelCase_         =     kwargs.pop("text"		,    lowercase					)
					if len(lowercase					) > 0:
									lowerCamelCase_         =     args[0]
									lowerCamelCase_         =     args[1:]

					if audio is None and text is None:
									raise ValueError("You need to specify either an `audio` or `text` input to process."					)

					if audio is not None:
									lowerCamelCase_         =     self.feature_extractor(lowercase		,    *lowercase		,    sampling_rate=lowercase		,    **lowercase					)
					if text is not None:
									lowerCamelCase_         =     self.tokenizer(lowercase		,    **lowercase					)

					if text is None:
									return inputs
					elif audio is None:
									return encodings
					else:
									lowerCamelCase_         =     encodings["input_ids"]
									return inputs



	def        SCREAMING_SNAKE_CASE_(       self		,    *lowercase		,    **lowercase					)							->			Optional[Any]:
					return self.tokenizer.batch_decode(*lowercase		,    **lowercase					)



	def        SCREAMING_SNAKE_CASE_(       self		,    *lowercase		,    **lowercase					)							->			Optional[Any]:
					# For backward compatibility
					if self._in_target_context_manager:
									return self.current_processor.pad(*lowercase		,    **lowercase					)

					lowerCamelCase_         =     kwargs.pop("input_features"		,    lowercase					)
					lowerCamelCase_         =     kwargs.pop("labels"		,    lowercase					)
					if len(lowercase					) > 0:
									lowerCamelCase_         =     args[0]
									lowerCamelCase_         =     args[1:]

					if input_features is not None:
									lowerCamelCase_         =     self.feature_extractor.pad(lowercase		,    *lowercase		,    **lowercase					)
					if labels is not None:
									lowerCamelCase_         =     self.tokenizer.pad(lowercase		,    **lowercase					)

					if labels is None:
									return input_features
					elif input_features is None:
									return labels
					else:
									lowerCamelCase_         =     labels["input_ids"]
									return input_features



	def        SCREAMING_SNAKE_CASE_(       self		,    *lowercase		,    **lowercase					)							->			List[str]:
					return self.tokenizer.decode(*lowercase		,    **lowercase					)



	@contextmanager
	def        SCREAMING_SNAKE_CASE_(       self					)							->			Optional[int]:
					warnings.warn(
					    "`as_target_processor` is deprecated and will be removed in v5 of Transformers. You can process your "
					    "labels by using the argument `text` of the regular `__call__` method (either in the same call as "
					    "your audio inputs, or in a separate call."					)
					lowerCamelCase_         =     True
					lowerCamelCase_         =     self.tokenizer
					yield
					lowerCamelCase_         =     self.feature_extractor
					lowerCamelCase_         =     False


