
"""simple docstring"""





from dataclasses import dataclass
from typing import List, Optional, Union

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

from ...models import PriorTransformer
from ...pipelines import DiffusionPipeline
from ...schedulers import HeunDiscreteScheduler
from ...utils import (
    BaseOutput,
    is_accelerate_available,
    logging,
    randn_tensor,
    replace_example_docstring,
)
from .renderer import ShapERenderer


_a														=		logging.get_logger(__name__)  # pylint: disable=invalid-name

_a														=		"""
    Examples:
        ```py
        >>> from PIL import Image
        >>> import torch
        >>> from diffusers import DiffusionPipeline
        >>> from diffusers.utils import export_to_gif, load_image

        >>> device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")

        >>> repo = \"openai/shap-e-img2img\"
        >>> pipe = DiffusionPipeline.from_pretrained(repo, torch_dtype=torch.float16)
        >>> pipe = pipe.to(device)

        >>> guidance_scale = 3.0
        >>> image_url = \"https://hf.co/datasets/diffusers/docs-images/resolve/main/shap-e/corgi.png\"
        >>> image = load_image(image_url).convert(\"RGB\")

        >>> images = pipe(
        ...     image,
        ...     guidance_scale=guidance_scale,
        ...     num_inference_steps=64,
        ...     frame_size=256,
        ... ).images

        >>> gif_path = export_to_gif(images[0], \"corgi_3d.gif\")
        ```
"""






@dataclass
class    _UpperCAmelCase(   lowerCamelCase     ):
		lowercase__      =				42






class    _UpperCAmelCase(   lowerCamelCase     ):



		def __init__( self    ,		__a    ,		__a    ,		__a    ,		__a    ,		__a    ,		)	->	Union[str, Any]:






							'''simple docstring'''
							super().__init__()

							self.register_modules(
							    prior=__a    ,		image_encoder=__a    ,		image_processor=__a    ,		scheduler=__a    ,		renderer=__a    ,		)



		def 						UpperCAmelCase							( self    ,		__a    ,		__a    ,		__a    ,		__a    ,		__a    ,		__a)	->	Union[str, Any]:






							'''simple docstring'''
							if latents is None:
												_UpperCamelCase										=  randn_tensor(__a    ,		generator=__a    ,		device=__a    ,		dtype=__a)
							else:
												if latents.shape != shape:
																	raise ValueError(F'''Unexpected latents shape, got {latents.shape}, expected {shape}''')
												_UpperCamelCase										=  latents.to(__a)

							_UpperCamelCase										=  latents * scheduler.init_noise_sigma
							return latents



		def 						UpperCAmelCase							( self    ,		__a=0)	->	Optional[Any]:






							'''simple docstring'''
							if is_accelerate_available():
												from accelerate import cpu_offload
							else:
												raise ImportError('''Please install accelerate via `pip install accelerate`''')

							_UpperCamelCase										=  torch.device(F'''cuda:{gpu_id}''')

							_UpperCamelCase										=  [self.image_encoder, self.prior]
							for cpu_offloaded_model in models:
												if cpu_offloaded_model is not None:
																	cpu_offload(__a    ,		__a)



		@property
		def 						UpperCAmelCase							( self)	->	List[str]:






							'''simple docstring'''
							if self.device != torch.device('''meta''') or not hasattr(self.image_encoder    ,		'''_hf_hook'''):
												return self.device
							for module in self.image_encoder.modules():
												if (
												    hasattr(__a    ,		'''_hf_hook''')
												    and hasattr(module._hf_hook    ,		'''execution_device''')
												    and module._hf_hook.execution_device is not None
												):
																	return torch.device(module._hf_hook.execution_device)
							return self.device



		def 						UpperCAmelCase							( self    ,		__a    ,		__a    ,		__a    ,		__a    ,		)	->	Optional[int]:






							'''simple docstring'''
							if isinstance(__a    ,		__a) and isinstance(image[0]    ,		torch.Tensor):
												_UpperCamelCase										=  torch.cat(__a    ,		axis=0) if image[0].ndim == 4 else torch.stack(__a    ,		axis=0)

							if not isinstance(__a    ,		torch.Tensor):
												_UpperCamelCase										=  self.image_processor(__a    ,		return_tensors='''pt''').pixel_values[0].unsqueeze(0)

							_UpperCamelCase										=  image.to(dtype=self.image_encoder.dtype    ,		device=__a)

							_UpperCamelCase										=  self.image_encoder(__a)['''last_hidden_state''']
							_UpperCamelCase										=  image_embeds[:, 1:, :].contiguous()  # batch_size, dim, 256

							_UpperCamelCase										=  image_embeds.repeat_interleave(__a    ,		dim=0)

							if do_classifier_free_guidance:
												_UpperCamelCase										=  torch.zeros_like(__a)

												# For classifier free guidance, we need to do two forward passes.
												# Here we concatenate the unconditional and text embeddings into a single batch
												# to avoid doing two forward passes
												_UpperCamelCase										=  torch.cat([negative_image_embeds, image_embeds])

							return image_embeds




		@torch.no_grad()
		@replace_example_docstring(__a)
		def __call__( self    ,		__a    ,		__a = 1    ,		__a = 25    ,		__a = None    ,		__a = None    ,		__a = 4.0    ,		__a = 64    ,		__a = "pil"    ,		__a = True    ,		)	->	Optional[int]:






							'''simple docstring'''

							if isinstance(__a    ,		PIL.Image.Image):
												_UpperCamelCase										=  1
							elif isinstance(__a    ,		torch.Tensor):
												_UpperCamelCase										=  image.shape[0]
							elif isinstance(__a    ,		__a) and isinstance(image[0]    ,		(torch.Tensor, PIL.Image.Image)):
												_UpperCamelCase										=  len(__a)
							else:
												raise ValueError(
												    F'''`image` has to be of type `PIL.Image.Image`, `torch.Tensor`, `List[PIL.Image.Image]` or `List[torch.Tensor]` but is {type(__a)}''')

							_UpperCamelCase										=  self._execution_device

							_UpperCamelCase										=  batch_size * num_images_per_prompt

							_UpperCamelCase										=  guidance_scale > 1.0
							_UpperCamelCase										=  self._encode_image(__a    ,		__a    ,		__a    ,		__a)

							# prior

							self.scheduler.set_timesteps(__a    ,		device=__a)
							_UpperCamelCase										=  self.scheduler.timesteps

							_UpperCamelCase										=  self.prior.config.num_embeddings
							_UpperCamelCase										=  self.prior.config.embedding_dim

							_UpperCamelCase										=  self.prepare_latents(
							    (batch_size, num_embeddings * embedding_dim)    ,		image_embeds.dtype    ,		__a    ,		__a    ,		__a    ,		self.scheduler    ,		)

							# YiYi notes: for testing only to match ldm, we can directly create a latents with desired shape: batch_size, num_embeddings, embedding_dim
							_UpperCamelCase										=  latents.reshape(latents.shape[0]    ,		__a    ,		__a)

							for i, t in enumerate(self.progress_bar(__a)):
												# expand the latents if we are doing classifier free guidance
												_UpperCamelCase										=  torch.cat([latents] * 2) if do_classifier_free_guidance else latents
												_UpperCamelCase										=  self.scheduler.scale_model_input(__a    ,		__a)

												_UpperCamelCase										=  self.prior(
												    __a    ,		timestep=__a    ,		proj_embedding=__a    ,		).predicted_image_embedding

												# remove the variance
												_UpperCamelCase			,  _UpperCamelCase										=  noise_pred.split(
												    scaled_model_input.shape[2]    ,		dim=2)  # batch_size, num_embeddings, embedding_dim

												if do_classifier_free_guidance is not None:
																	_UpperCamelCase			,  _UpperCamelCase										=  noise_pred.chunk(2)
																	_UpperCamelCase										=  noise_pred_uncond + guidance_scale * (noise_pred - noise_pred_uncond)

												_UpperCamelCase										=  self.scheduler.step(
												    __a    ,		timestep=__a    ,		sample=__a    ,		).prev_sample

							if output_type == "latent":
												return ShapEPipelineOutput(images=__a)

							_UpperCamelCase										=  []
							for i, latent in enumerate(__a):
												print()
												_UpperCamelCase										=  self.renderer.decode(
												    latent[None, :]    ,		__a    ,		size=__a    ,		ray_batch_size=40_96    ,		n_coarse_samples=64    ,		n_fine_samples=1_28    ,		)

												images.append(__a)

							_UpperCamelCase										=  torch.stack(__a)

							if output_type not in ["np", "pil"]:
												raise ValueError(F'''Only the output types `pil` and `np` are supported not output_type={output_type}''')

							_UpperCamelCase										=  images.cpu().numpy()

							if output_type == "pil":
												_UpperCamelCase										=  [self.numpy_to_pil(__a) for image in images]

							# Offload last model to CPU
							if hasattr(self    ,		'''final_offload_hook''') and self.final_offload_hook is not None:
												self.final_offload_hook.offload()

							if not return_dict:
												return (images,)

							return ShapEPipelineOutput(images=__a)








"""simple docstring"""





import importlib
import json
import os
import sys
import tempfile
import unittest
from pathlib import Path

import transformers
import transformers.models.auto
from transformers.models.auto.configuration_auto import CONFIG_MAPPING, AutoConfig
from transformers.models.bert.configuration_bert import BertConfig
from transformers.models.roberta.configuration_roberta import RobertaConfig
from transformers.testing_utils import DUMMY_UNKNOWN_IDENTIFIER, get_tests_dir


sys.path.append(str(Path(__file__).parent.parent.parent.parent / """utils"""))

from test_module.custom_configuration import CustomConfig  # noqa E402


_a														=		get_tests_dir("""fixtures/dummy-config.json""")






class    _UpperCAmelCase(   unittest.TestCase     ):



		def 						UpperCAmelCase							( self)	->	str:






							'''simple docstring'''
							_UpperCamelCase										=  0



		def 						UpperCAmelCase							( self)	->	List[str]:






							'''simple docstring'''
							self.assertIsNotNone(transformers.models.auto.__spec__)
							self.assertIsNotNone(importlib.util.find_spec('''transformers.models.auto'''))



		def 						UpperCAmelCase							( self)	->	str:






							'''simple docstring'''
							_UpperCamelCase										=  AutoConfig.from_pretrained('''bert-base-uncased''')
							self.assertIsInstance(__a    ,		__a)



		def 						UpperCAmelCase							( self)	->	Tuple:






							'''simple docstring'''
							_UpperCamelCase										=  AutoConfig.from_pretrained(__a)
							self.assertIsInstance(__a    ,		__a)



		def 						UpperCAmelCase							( self)	->	Optional[Any]:






							'''simple docstring'''
							_UpperCamelCase										=  AutoConfig.from_pretrained(__a)
							self.assertIsInstance(__a    ,		__a)



		def 						UpperCAmelCase							( self)	->	Union[str, Any]:






							'''simple docstring'''
							_UpperCamelCase										=  AutoConfig.for_model('''roberta''')
							self.assertIsInstance(__a    ,		__a)



		def 						UpperCAmelCase							( self)	->	Dict:






							'''simple docstring'''
							with tempfile.TemporaryDirectory() as tmp_dir:
												# This model name contains bert and roberta, but roberta ends up being picked.
												_UpperCamelCase										=  os.path.join(__a    ,		'''fake-roberta''')
												os.makedirs(__a    ,		exist_ok=__a)
												with open(os.path.join(__a    ,		'''config.json''')    ,		'''w''') as f:
																	f.write(json.dumps({}))
												_UpperCamelCase										=  AutoConfig.from_pretrained(__a)
												self.assertEqual(type(__a)    ,		__a)



		def 						UpperCAmelCase							( self)	->	Tuple:






							'''simple docstring'''
							try:
												AutoConfig.register('''custom'''    ,		__a)
												# Wrong model type will raise an error
												with self.assertRaises(__a):
																	AutoConfig.register('''model'''    ,		__a)
												# Trying to register something existing in the Transformers library will raise an error
												with self.assertRaises(__a):
																	AutoConfig.register('''bert'''    ,		__a)

												# Now that the config is registered, it can be used as any other config with the auto-API
												_UpperCamelCase										=  CustomConfig()
												with tempfile.TemporaryDirectory() as tmp_dir:
																	config.save_pretrained(__a)
																	_UpperCamelCase										=  AutoConfig.from_pretrained(__a)
																	self.assertIsInstance(__a    ,		__a)

							finally:
												if "custom" in CONFIG_MAPPING._extra_content:
																	del CONFIG_MAPPING._extra_content["custom"]



		def 						UpperCAmelCase							( self)	->	List[Any]:






							'''simple docstring'''
							with self.assertRaisesRegex(
							    __a    ,		'''bert-base is not a local folder and is not a valid model identifier'''):
												_UpperCamelCase										=  AutoConfig.from_pretrained('''bert-base''')



		def 						UpperCAmelCase							( self)	->	List[Any]:






							'''simple docstring'''
							with self.assertRaisesRegex(
							    __a    ,		R'''aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)'''):
												_UpperCamelCase										=  AutoConfig.from_pretrained(__a    ,		revision='''aaaaaa''')



		def 						UpperCAmelCase							( self)	->	Optional[Any]:






							'''simple docstring'''
							with self.assertRaisesRegex(
							    __a    ,		'''hf-internal-testing/no-config-test-repo does not appear to have a file named config.json.'''    ,		):
												_UpperCamelCase										=  AutoConfig.from_pretrained('''hf-internal-testing/no-config-test-repo''')



		def 						UpperCAmelCase							( self)	->	int:






							'''simple docstring'''
							# If remote code is not set, we will time out when asking whether to load the model.
							with self.assertRaises(__a):
												_UpperCamelCase										=  AutoConfig.from_pretrained('''hf-internal-testing/test_dynamic_model''')
							# If remote code is disabled, we can't load this config.
							with self.assertRaises(__a):
												_UpperCamelCase										=  AutoConfig.from_pretrained('''hf-internal-testing/test_dynamic_model'''    ,		trust_remote_code=__a)

							_UpperCamelCase										=  AutoConfig.from_pretrained('''hf-internal-testing/test_dynamic_model'''    ,		trust_remote_code=__a)
							self.assertEqual(config.__class__.__name__    ,		'''NewModelConfig''')

							# Test config can be reloaded.
							with tempfile.TemporaryDirectory() as tmp_dir:
												config.save_pretrained(__a)
												_UpperCamelCase										=  AutoConfig.from_pretrained(__a    ,		trust_remote_code=__a)
							self.assertEqual(reloaded_config.__class__.__name__    ,		'''NewModelConfig''')




		def 						UpperCAmelCase							( self)	->	Any:






							'''simple docstring'''






							class    _UpperCAmelCase(   lowerCamelCase     ):
									lowercase__      =				'new-model'

							try:
												AutoConfig.register('''new-model'''    ,		__a)
												# If remote code is not set, the default is to use local
												_UpperCamelCase										=  AutoConfig.from_pretrained('''hf-internal-testing/test_dynamic_model''')
												self.assertEqual(config.__class__.__name__    ,		'''NewModelConfigLocal''')

												# If remote code is disabled, we load the local one.
												_UpperCamelCase										=  AutoConfig.from_pretrained('''hf-internal-testing/test_dynamic_model'''    ,		trust_remote_code=__a)
												self.assertEqual(config.__class__.__name__    ,		'''NewModelConfigLocal''')

												# If remote is enabled, we load from the Hub
												_UpperCamelCase										=  AutoConfig.from_pretrained('''hf-internal-testing/test_dynamic_model'''    ,		trust_remote_code=__a)
												self.assertEqual(config.__class__.__name__    ,		'''NewModelConfig''')

							finally:
												if "new-model" in CONFIG_MAPPING._extra_content:
																	del CONFIG_MAPPING._extra_content["new-model"]






