





import os
import sys
import warnings
from dataclasses import dataclass, field
from io import BytesIO
from typing import TYPE_CHECKING, Any, ClassVar, Dict, List, Optional, Union

import numpy as np
import pyarrow as pa

from .. import config
from ..download.streaming_download_manager import xopen
from ..table import array_cast
from ..utils.file_utils import is_local_path
from ..utils.py_utils import first_non_null_value, no_op_if_value_is_null, string_to_dict


if TYPE_CHECKING:
	import PIL.Image

	from .features import FeatureType


_lowerCamelCase 					=      None
_lowerCamelCase 					=      '<' if sys.byteorder == 'little' else '>'
# Origin: https://github.com/python-pillow/Pillow/blob/698951e19e19972aeed56df686868f1329981c12/src/PIL/Image.py#L3126 minus "|i1" which values are not preserved correctly when saving and loading an image
_lowerCamelCase 					=      [
    np.dtype('|b1'),
    np.dtype('|u1'),
    np.dtype('<u2'),
    np.dtype('>u2'),
    np.dtype('<i2'),
    np.dtype('>i2'),
    np.dtype('<u4'),
    np.dtype('>u4'),
    np.dtype('<i4'),
    np.dtype('>i4'),
    np.dtype('<f4'),
    np.dtype('>f4'),
    np.dtype('<f8'),
    np.dtype('>f8'),
]




@dataclass
class 	UpperCamelCase_							:
							lowerCamelCase_							     =       True
							lowerCamelCase_							     =       None
							# Automatically constructed
							lowerCamelCase_							     =       "PIL.Image.Image"
							lowerCamelCase_							     =       pa.struct({"bytes": pa.binary(), "path": pa.string()}     )
							lowerCamelCase_							     =       field(default="Image"					,	init=UpperCamelCase__					,	repr=UpperCamelCase__     )


							def __call__(				self     :Optional[Any]		) -> Optional[int]:



									"""simple docstring"""



									return self.pa_type


							def   _snake_case	(				self     :int				,					__A     :Union[str, bytes, dict, np.ndarray, "PIL.Image.Image"]		) -> dict:



									"""simple docstring"""



									if config.PIL_AVAILABLE:
											import PIL.Image
									else:
											raise ImportError("""To support encoding images, please install 'Pillow'."""		)

									if isinstance(__A				,					__A		):
											SCREAMING_SNAKE_CASE__								=   np.array(__A		)

									if isinstance(__A				,					__A		):
											return {"path": value, "bytes": None}
									elif isinstance(__A				,					__A		):
											return {"path": None, "bytes": value}
									elif isinstance(__A				,					np.ndarray		):
											# convert the image array to PNG/TIFF bytes
											return encode_np_array(__A		)
									elif isinstance(__A				,					PIL.Image.Image		):
											# convert the PIL image to bytes (default format is PNG/TIFF)
											return encode_pil_image(__A		)
									elif value.get("""path"""		) is not None and os.path.isfile(value["""path"""]		):
											# we set "bytes": None to not duplicate the data if they're already available locally
											return {"bytes": None, "path": value.get("""path"""		)}
									elif value.get("""bytes"""		) is not None or value.get("""path"""		) is not None:
											# store the image bytes, and path is used to infer the image format using the file extension
											return {"bytes": value.get("""bytes"""		), "path": value.get("""path"""		)}
									else:
											raise ValueError(
											    f'''An image sample should have one of \'path\' or \'bytes\' but they are missing or None in {value}.'''		)


							def   _snake_case	(				self     :Tuple				,					__A     :dict				,					__A     :Optional[Any]=None		) -> "PIL.Image.Image":



									"""simple docstring"""



									if not self.decode:
											raise RuntimeError("""Decoding is disabled for this feature. Please use Image(decode=True) instead."""		)

									if config.PIL_AVAILABLE:
											import PIL.Image
									else:
											raise ImportError("""To support decoding images, please install 'Pillow'."""		)

									if token_per_repo_id is None:
											SCREAMING_SNAKE_CASE__								=   {}

									SCREAMING_SNAKE_CASE__						,    SCREAMING_SNAKE_CASE__								=   value["""path"""], value["""bytes"""]
									if bytes_ is None:
											if path is None:
													raise ValueError(f'''An image should have one of \'path\' or \'bytes\' but both are None in {value}.'''		)
											else:
													if is_local_path(__A		):
															SCREAMING_SNAKE_CASE__								=   PIL.Image.open(__A		)
													else:
															SCREAMING_SNAKE_CASE__								=   path.split("""::"""		)[-1]
															try:
																	SCREAMING_SNAKE_CASE__								=   string_to_dict(__A				,					config.HUB_DATASETS_URL		)["""repo_id"""]
																	SCREAMING_SNAKE_CASE__								=   token_per_repo_id.get(__A		)
															except ValueError:
																	SCREAMING_SNAKE_CASE__								=   None
															with xopen(__A				,					"""rb"""				,					use_auth_token=__A		) as f:
																	SCREAMING_SNAKE_CASE__								=   BytesIO(f.read()		)
															SCREAMING_SNAKE_CASE__								=   PIL.Image.open(bytes_		)
									else:
											SCREAMING_SNAKE_CASE__								=   PIL.Image.open(BytesIO(bytes_		)		)
									image.load()  # to avoid "Too many open files" errors
									return image


							def   _snake_case	(				self     :str		) -> Union["FeatureType", Dict[str, "FeatureType"]]:



									"""simple docstring"""



									from .features import Value

									return (
									    self
									    if self.decode
									    else {
									        "bytes": Value("""binary"""		),
									        "path": Value("""string"""		),
									    }
									)


							def   _snake_case	(				self     :List[Any]				,					__A     :Union[pa.StringArray, pa.StructArray, pa.ListArray]		) -> pa.StructArray:



									"""simple docstring"""



									if pa.types.is_string(storage.type		):
											SCREAMING_SNAKE_CASE__								=   pa.array([None] * len(__A		)				,					type=pa.binary()		)
											SCREAMING_SNAKE_CASE__								=   pa.StructArray.from_arrays([bytes_array, storage]				,					["""bytes""", """path"""]				,					mask=storage.is_null()		)
									elif pa.types.is_binary(storage.type		):
											SCREAMING_SNAKE_CASE__								=   pa.array([None] * len(__A		)				,					type=pa.string()		)
											SCREAMING_SNAKE_CASE__								=   pa.StructArray.from_arrays([storage, path_array]				,					["""bytes""", """path"""]				,					mask=storage.is_null()		)
									elif pa.types.is_struct(storage.type		):
											if storage.type.get_field_index("""bytes"""		) >= 0:
													SCREAMING_SNAKE_CASE__								=   storage.field("""bytes"""		)
											else:
													SCREAMING_SNAKE_CASE__								=   pa.array([None] * len(__A		)				,					type=pa.binary()		)
											if storage.type.get_field_index("""path"""		) >= 0:
													SCREAMING_SNAKE_CASE__								=   storage.field("""path"""		)
											else:
													SCREAMING_SNAKE_CASE__								=   pa.array([None] * len(__A		)				,					type=pa.string()		)
											SCREAMING_SNAKE_CASE__								=   pa.StructArray.from_arrays([bytes_array, path_array]				,					["""bytes""", """path"""]				,					mask=storage.is_null()		)
									elif pa.types.is_list(storage.type		):
											SCREAMING_SNAKE_CASE__								=   pa.array(
											    [encode_np_array(np.array(__A		)		)["""bytes"""] if arr is not None else None for arr in storage.to_pylist()]				,					type=pa.binary()				,					)
											SCREAMING_SNAKE_CASE__								=   pa.array([None] * len(__A		)				,					type=pa.string()		)
											SCREAMING_SNAKE_CASE__								=   pa.StructArray.from_arrays(
											    [bytes_array, path_array]				,					["""bytes""", """path"""]				,					mask=bytes_array.is_null()		)
									return array_cast(__A				,					self.pa_type		)


							def   _snake_case	(				self     :Tuple				,					__A     :pa.StructArray		) -> pa.StructArray:



									"""simple docstring"""



									@no_op_if_value_is_null
									def path_to_bytes(__A     :Tuple		):
											with xopen(__A				,					"""rb"""		) as f:
													SCREAMING_SNAKE_CASE__								=   f.read()
											return bytes_

									SCREAMING_SNAKE_CASE__								=   pa.array(
									    [
									        (path_to_bytes(x["""path"""]		) if x["""bytes"""] is None else x["""bytes"""]) if x is not None else None
									        for x in storage.to_pylist()
									    ]				,					type=pa.binary()				,					)
									SCREAMING_SNAKE_CASE__								=   pa.array(
									    [os.path.basename(__A		) if path is not None else None for path in storage.field("""path"""		).to_pylist()]				,					type=pa.string()				,					)
									SCREAMING_SNAKE_CASE__								=   pa.StructArray.from_arrays([bytes_array, path_array]				,					["""bytes""", """path"""]				,					mask=bytes_array.is_null()		)
									return array_cast(__A				,					self.pa_type		)






def 					SCREAMING_SNAKE_CASE__				(							):
		if config.PIL_AVAILABLE:
				import PIL.Image
		else:
				raise ImportError("""To support encoding images, please install 'Pillow'."""     )

		global _IMAGE_COMPRESSION_FORMATS
		if _IMAGE_COMPRESSION_FORMATS is None:
				PIL.Image.init()
				SCREAMING_SNAKE_CASE__								=   list(set(PIL.Image.OPEN.keys()     ) & set(PIL.Image.SAVE.keys()     )     )
		return _IMAGE_COMPRESSION_FORMATS






def 					SCREAMING_SNAKE_CASE__				(							UpperCamelCase__:				"PIL.Image.Image"     ):
		SCREAMING_SNAKE_CASE__								=   BytesIO()
		if image.format in list_image_compression_formats():
				SCREAMING_SNAKE_CASE__								=   image.format
		else:
				SCREAMING_SNAKE_CASE__								=   """PNG""" if image.mode in ["""1""", """L""", """LA""", """RGB""", """RGBA"""] else """TIFF"""
		image.save(UpperCamelCase__       , format=UpperCamelCase__     )
		return buffer.getvalue()






def 					SCREAMING_SNAKE_CASE__				(							UpperCamelCase__:				"PIL.Image.Image"     ):
		if hasattr(UpperCamelCase__       , """filename"""     ) and image.filename != "":
				return {"path": image.filename, "bytes": None}
		else:
				return {"path": None, "bytes": image_to_bytes(UpperCamelCase__     )}






def 					SCREAMING_SNAKE_CASE__				(							UpperCamelCase__:				np.ndarray     ):
		if config.PIL_AVAILABLE:
				import PIL.Image
		else:
				raise ImportError("""To support encoding images, please install 'Pillow'."""     )

		SCREAMING_SNAKE_CASE__								=   array.dtype
		SCREAMING_SNAKE_CASE__								=   dtype.byteorder if dtype.byteorder != """=""" else _NATIVE_BYTEORDER
		SCREAMING_SNAKE_CASE__								=   dtype.kind
		SCREAMING_SNAKE_CASE__								=   dtype.itemsize

		SCREAMING_SNAKE_CASE__								=   None

		# Multi-channel array case (only np.dtype("|u1") is allowed)
		if array.shape[2:]:
				SCREAMING_SNAKE_CASE__								=   np.dtype("""|u1"""     )
				if dtype_kind not in ["u", "i"]:
						raise TypeError(
						    f'''Unsupported array dtype {dtype} for image encoding. Only {dest_dtype} is supported for multi-channel arrays.'''     )
				if dtype is not dest_dtype:
						warnings.warn(f'''Downcasting array dtype {dtype} to {dest_dtype} to be compatible with \'Pillow\''''     )
    # Exact match
		elif dtype in _VALID_IMAGE_ARRAY_DTPYES:
				SCREAMING_SNAKE_CASE__								=   dtype
		else:  # Downcast the type within the kind (np.can_cast(from_type, to_type, casting="same_kind") doesn't behave as expected, so do it manually)
				while dtype_itemsize >= 1:
						SCREAMING_SNAKE_CASE__								=   dtype_byteorder + dtype_kind + str(UpperCamelCase__     )
						SCREAMING_SNAKE_CASE__								=   np.dtype(UpperCamelCase__     )
						if dest_dtype in _VALID_IMAGE_ARRAY_DTPYES:
								warnings.warn(f'''Downcasting array dtype {dtype} to {dest_dtype} to be compatible with \'Pillow\''''     )
								break
						else:
								dtype_itemsize //= 2

		if dest_dtype is None:
				raise TypeError(
				    f'''Cannot convert dtype {dtype} to a valid image dtype. Valid image dtypes: {_VALID_IMAGE_ARRAY_DTPYES}'''     )
		SCREAMING_SNAKE_CASE__								=   PIL.Image.fromarray(array.astype(UpperCamelCase__     )     )
		return {"path": None, "bytes": image_to_bytes(UpperCamelCase__     )}






def 					SCREAMING_SNAKE_CASE__				(							UpperCamelCase__:				Union[List[str], List[dict], List[np.ndarray], List["PIL.Image.Image"]]     ):
		if config.PIL_AVAILABLE:
				import PIL.Image
		else:
				raise ImportError("""To support encoding images, please install 'Pillow'."""     )

		if objs:
				SCREAMING_SNAKE_CASE__						,    SCREAMING_SNAKE_CASE__								=   first_non_null_value(UpperCamelCase__     )
				if isinstance(UpperCamelCase__       , UpperCamelCase__     ):
						return [{"path": obj, "bytes": None} if obj is not None else None for obj in objs]
				if isinstance(UpperCamelCase__       , np.ndarray     ):
						SCREAMING_SNAKE_CASE__								=   no_op_if_value_is_null(UpperCamelCase__     )
						return [obj_to_image_dict_func(UpperCamelCase__     ) for obj in objs]
				elif isinstance(UpperCamelCase__       , PIL.Image.Image     ):
						SCREAMING_SNAKE_CASE__								=   no_op_if_value_is_null(UpperCamelCase__     )
						return [obj_to_image_dict_func(UpperCamelCase__     ) for obj in objs]
				else:
						return objs
		else:
				return objs
'''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


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

__lowercase        =    '''
    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 							a__(   lowerCAmelCase__      ):


			'''simple docstring'''






			UpperCAmelCase_     :	Union[PIL.Image.Image, np.ndarray]







class 							a__(   lowerCAmelCase__      ):


			'''simple docstring'''






			def __init__(  self			,       __lowerCAmelCase			,       __lowerCAmelCase			,       __lowerCAmelCase			,       __lowerCAmelCase			,       __lowerCAmelCase			,       ):




					"""simple docstring"""


					super().__init__()

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



			def 				a_			(  self			,       __lowerCAmelCase			,       __lowerCAmelCase			,       __lowerCAmelCase			,       __lowerCAmelCase			,       __lowerCAmelCase			,       __lowerCAmelCase):




					"""simple docstring"""


					if latents is None:
							lowerCAmelCase     =    randn_tensor(__lowerCAmelCase			,       generator=__lowerCAmelCase			,       device=__lowerCAmelCase			,       dtype=__lowerCAmelCase)
					else:
							if latents.shape != shape:
									raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {shape}")
							lowerCAmelCase     =    latents.to(__lowerCAmelCase)

					lowerCAmelCase     =    latents * scheduler.init_noise_sigma
					return latents



			def 				a_			(  self			,       __lowerCAmelCase=0):




					"""simple docstring"""


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

					lowerCAmelCase     =    torch.device(f"cuda:{gpu_id}")

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



			@property
			def 				a_			(  self):




					"""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(__lowerCAmelCase			,       """_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 				a_			(  self			,       __lowerCAmelCase			,       __lowerCAmelCase			,       __lowerCAmelCase			,       __lowerCAmelCase			,       ):




					"""simple docstring"""


					if isinstance(__lowerCAmelCase			,       __lowerCAmelCase) and isinstance(image[0]			,       torch.Tensor):
							lowerCAmelCase     =    torch.cat(__lowerCAmelCase			,       axis=0) if image[0].ndim == 4 else torch.stack(__lowerCAmelCase			,       axis=0)

					if not isinstance(__lowerCAmelCase			,       torch.Tensor):
							lowerCAmelCase     =    self.image_processor(__lowerCAmelCase			,       return_tensors="""pt""").pixel_values[0].unsqueeze(0)

					lowerCAmelCase     =    image.to(dtype=self.image_encoder.dtype			,       device=__lowerCAmelCase)

					lowerCAmelCase     =    self.image_encoder(__lowerCAmelCase)["""last_hidden_state"""]
					lowerCAmelCase     =    image_embeds[:, 1:, :].contiguous()  # batch_size, dim, 256

					lowerCAmelCase     =    image_embeds.repeat_interleave(__lowerCAmelCase			,       dim=0)

					if do_classifier_free_guidance:
							lowerCAmelCase     =    torch.zeros_like(__lowerCAmelCase)

							# 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
							lowerCAmelCase     =    torch.cat([negative_image_embeds, image_embeds])

					return image_embeds







			@torch.no_grad()
			@replace_example_docstring(__lowerCAmelCase)
			def __call__(  self			,       __lowerCAmelCase			,       __lowerCAmelCase = 1			,       __lowerCAmelCase = 25			,       __lowerCAmelCase = None			,       __lowerCAmelCase = None			,       __lowerCAmelCase = 4.0			,       __lowerCAmelCase = 64			,       __lowerCAmelCase = "pil"			,       __lowerCAmelCase = True			,       ):




					"""simple docstring"""


					if isinstance(__lowerCAmelCase			,       PIL.Image.Image):
							lowerCAmelCase     =    1
					elif isinstance(__lowerCAmelCase			,       torch.Tensor):
							lowerCAmelCase     =    image.shape[0]
					elif isinstance(__lowerCAmelCase			,       __lowerCAmelCase) and isinstance(image[0]			,       (torch.Tensor, PIL.Image.Image)):
							lowerCAmelCase     =    len(__lowerCAmelCase)
					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(__lowerCAmelCase)}")

					lowerCAmelCase     =    self._execution_device

					lowerCAmelCase     =    batch_size * num_images_per_prompt

					lowerCAmelCase     =    guidance_scale > 1.0
					lowerCAmelCase     =    self._encode_image(__lowerCAmelCase			,       __lowerCAmelCase			,       __lowerCAmelCase			,       __lowerCAmelCase)

					# prior

					self.scheduler.set_timesteps(__lowerCAmelCase			,       device=__lowerCAmelCase)
					lowerCAmelCase     =    self.scheduler.timesteps

					lowerCAmelCase     =    self.prior.config.num_embeddings
					lowerCAmelCase     =    self.prior.config.embedding_dim

					lowerCAmelCase     =    self.prepare_latents(
					    (batch_size, num_embeddings * embedding_dim)			,       image_embeds.dtype			,       __lowerCAmelCase			,       __lowerCAmelCase			,       __lowerCAmelCase			,       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
					lowerCAmelCase     =    latents.reshape(latents.shape[0]			,       __lowerCAmelCase			,       __lowerCAmelCase)

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

							lowerCAmelCase     =    self.prior(
							    __lowerCAmelCase			,       timestep=__lowerCAmelCase			,       proj_embedding=__lowerCAmelCase			,       ).predicted_image_embedding

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

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

							lowerCAmelCase     =    self.scheduler.step(
							    __lowerCAmelCase			,       timestep=__lowerCAmelCase			,       sample=__lowerCAmelCase			,       ).prev_sample

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

					lowerCAmelCase     =    []
					for i, latent in enumerate(__lowerCAmelCase):
							print()
							lowerCAmelCase     =    self.renderer.decode(
							    latent[None, :]			,       __lowerCAmelCase			,       size=__lowerCAmelCase			,       ray_batch_size=4096			,       n_coarse_samples=64			,       n_fine_samples=128			,       )

							images.append(__lowerCAmelCase)

					lowerCAmelCase     =    torch.stack(__lowerCAmelCase)

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

					lowerCAmelCase     =    images.cpu().numpy()

					if output_type == "pil":
							lowerCAmelCase     =    [self.numpy_to_pil(__lowerCAmelCase) 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=__lowerCAmelCase)




