



from typing import Optional, Tuple, Union

import tensorflow as tf

from ...activations_tf import ACTaFN
from ...modeling_tf_outputs import (
    TFBaseModelOutputWithNoAttention,
    TFBaseModelOutputWithPoolingAndNoAttention,
    TFImageClassifierOutputWithNoAttention,
)
from ...modeling_tf_utils import TFPreTrainedModel, TFSequenceClassificationLoss, keras_serializable, unpack_inputs
from ...tf_utils import shape_list
from ...utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward, logging
from .configuration_resnet import ResNetConfig


UpperCAmelCase							=       logging.get_logger(__name__)

# General docstring
UpperCAmelCase							=       """ResNetConfig"""

# Base docstring
UpperCAmelCase							=       """microsoft/resnet-50"""
UpperCAmelCase							=       [1, 2_048, 7, 7]

# Image classification docstring
UpperCAmelCase							=       """microsoft/resnet-50"""
UpperCAmelCase							=       """tiger cat"""

UpperCAmelCase							=       [
    """microsoft/resnet-50""",
    # See all resnet models at https://huggingface.co/models?filter=resnet
]

class 			__snake_case   ( tf.keras.layers.Layer):







							'''simple docstring'''






							def __init__(				self	,	a_	,	a_ = 3	,	a_ = 1	,	a_ = "relu"	,	**a_		):
											super().__init__(**a_		)
											a__					   =		kernel_size // 2
											a__					   =		tf.keras.layers.ConvaD(
											    a_	,	kernel_size=a_	,	strides=a_	,	padding="""valid"""	,	use_bias=a_	,	name="""convolution"""		)
											# Use same default momentum and epsilon as PyTorch equivalent
											a__					   =		tf.keras.layers.BatchNormalization(epsilon=1E-5	,	momentum=0.9	,	name="""normalization"""		)
											a__					   =		ACTaFN[activation] if activation is not None else tf.keras.layers.Activation("""linear"""		)


							def _a   (				self	,	a_		):
											# Pad to match that done in the PyTorch Conv2D model
											a__					   =		a__					   =		(self.pad_value, self.pad_value)
											a__					   =		tf.pad(a_	,	[(0, 0), height_pad, width_pad, (0, 0)]		)
											a__					   =		self.conv(a_		)
											return hidden_state







							def _a   (				self	,	a_	,	a_ = False		):
											a__					   =		self.convolution(a_		)
											a__					   =		self.normalization(a_	,	training=a_		)
											a__					   =		self.activation(a_		)
											return hidden_state




class 			__snake_case   ( tf.keras.layers.Layer):







							'''simple docstring'''






							def __init__(				self	,	a_	,	**a_		):
											super().__init__(**a_		)
											a__					   =		TFResNetConvLayer(
											    config.embedding_size	,	kernel_size=7	,	stride=2	,	activation=config.hidden_act	,	name="""embedder"""	,	)
											a__					   =		tf.keras.layers.MaxPoolaD(pool_size=3	,	strides=2	,	padding="""valid"""	,	name="""pooler"""		)
											a__					   =		config.num_channels







							def _a   (				self	,	a_	,	a_ = False		):
											a__   ,      a__   ,      a__   ,      a__					   =		shape_list(a_		)
											if tf.executing_eagerly() and num_channels != self.num_channels:
															raise ValueError(
															    """Make sure that the channel dimension of the pixel values match with the one set in the configuration."""		)
											a__					   =		pixel_values
											a__					   =		self.embedder(a_		)
											a__					   =		tf.pad(a_	,	[[0, 0], [1, 1], [1, 1], [0, 0]]		)
											a__					   =		self.pooler(a_		)
											return hidden_state




class 			__snake_case   ( tf.keras.layers.Layer):







							'''simple docstring'''






							def __init__(				self	,	a_	,	a_ = 2	,	**a_		):
											super().__init__(**a_		)
											a__					   =		tf.keras.layers.ConvaD(
											    a_	,	kernel_size=1	,	strides=a_	,	use_bias=a_	,	name="""convolution"""		)
											# Use same default momentum and epsilon as PyTorch equivalent
											a__					   =		tf.keras.layers.BatchNormalization(epsilon=1E-5	,	momentum=0.9	,	name="""normalization"""		)







							def _a   (				self	,	a_	,	a_ = False		):
											a__					   =		x
											a__					   =		self.convolution(a_		)
											a__					   =		self.normalization(a_	,	training=a_		)
											return hidden_state




class 			__snake_case   ( tf.keras.layers.Layer):







							'''simple docstring'''






							def __init__(				self	,	a_	,	a_	,	a_ = 1	,	a_ = "relu"	,	**a_		):
											super().__init__(**a_		)
											a__					   =		in_channels != out_channels or stride != 1
											a__					   =		TFResNetConvLayer(a_	,	stride=a_	,	name="""layer.0"""		)
											a__					   =		TFResNetConvLayer(a_	,	activation=a_	,	name="""layer.1"""		)
											a__					   =		(
											    TFResNetShortCut(a_	,	stride=a_	,	name="""shortcut"""		)
											    if should_apply_shortcut
											    else tf.keras.layers.Activation("""linear"""	,	name="""shortcut"""		)
											)
											a__					   =		ACTaFN[activation]







							def _a   (				self	,	a_	,	a_ = False		):
											a__					   =		hidden_state
											a__					   =		self.conva(a_	,	training=a_		)
											a__					   =		self.conva(a_	,	training=a_		)
											a__					   =		self.shortcut(a_	,	training=a_		)
											hidden_state += residual
											a__					   =		self.activation(a_		)
											return hidden_state




class 			__snake_case   ( tf.keras.layers.Layer):







							'''simple docstring'''






							def __init__(				self	,	a_	,	a_	,	a_ = 1	,	a_ = "relu"	,	a_ = 4	,	**a_	,	):
											super().__init__(**a_		)
											a__					   =		in_channels != out_channels or stride != 1
											a__					   =		out_channels // reduction
											a__					   =		TFResNetConvLayer(a_	,	kernel_size=1	,	name="""layer.0"""		)
											a__					   =		TFResNetConvLayer(a_	,	stride=a_	,	name="""layer.1"""		)
											a__					   =		TFResNetConvLayer(a_	,	kernel_size=1	,	activation=a_	,	name="""layer.2"""		)
											a__					   =		(
											    TFResNetShortCut(a_	,	stride=a_	,	name="""shortcut"""		)
											    if should_apply_shortcut
											    else tf.keras.layers.Activation("""linear"""	,	name="""shortcut"""		)
											)
											a__					   =		ACTaFN[activation]







							def _a   (				self	,	a_	,	a_ = False		):
											a__					   =		hidden_state
											a__					   =		self.conva(a_	,	training=a_		)
											a__					   =		self.conva(a_	,	training=a_		)
											a__					   =		self.conva(a_	,	training=a_		)
											a__					   =		self.shortcut(a_	,	training=a_		)
											hidden_state += residual
											a__					   =		self.activation(a_		)
											return hidden_state




class 			__snake_case   ( tf.keras.layers.Layer):







							'''simple docstring'''






							def __init__(				self	,	a_	,	a_	,	a_	,	a_ = 2	,	a_ = 2	,	**a_		):
											super().__init__(**a_		)

											a__					   =		TFResNetBottleNeckLayer if config.layer_type == """bottleneck""" else TFResNetBasicLayer

											a__					   =		[layer(a_	,	a_	,	stride=a_	,	activation=config.hidden_act	,	name="""layers.0"""		)]
											layers += [
											    layer(a_	,	a_	,	activation=config.hidden_act	,	name=F'''layers.{i + 1}'''		)
											    for i in range(depth - 1		)
											]
											a__					   =		layers







							def _a   (				self	,	a_	,	a_ = False		):
											for layer in self.stage_layers:
															a__					   =		layer(a_	,	training=a_		)
											return hidden_state




class 			__snake_case   ( tf.keras.layers.Layer):







							'''simple docstring'''






							def __init__(				self	,	a_	,	**a_		):
											super().__init__(**a_		)
											# based on `downsample_in_first_stage` the first layer of the first stage may or may not downsample the input
											a__					   =		[
											    TFResNetStage(
											        a_	,	config.embedding_size	,	config.hidden_sizes[0]	,	stride=2 if config.downsample_in_first_stage else 1	,	depth=config.depths[0]	,	name="""stages.0"""	,	)
											]
											for i, (in_channels, out_channels, depth) in enumerate(
											    zip(config.hidden_sizes	,	config.hidden_sizes[1:]	,	config.depths[1:]		)		):
															self.stages.append(TFResNetStage(a_	,	a_	,	a_	,	depth=a_	,	name=F'''stages.{i + 1}'''		)		)







							def _a   (				self	,	a_	,	a_ = False	,	a_ = True	,	a_ = False	,	):
											a__					   =		() if output_hidden_states else None

											for stage_module in self.stages:
															if output_hidden_states:
																			a__					   =		hidden_states + (hidden_state,)

															a__					   =		stage_module(a_	,	training=a_		)

											if output_hidden_states:
															a__					   =		hidden_states + (hidden_state,)

											if not return_dict:
															return tuple(v for v in [hidden_state, hidden_states] if v is not None		)

											return TFBaseModelOutputWithNoAttention(last_hidden_state=a_	,	hidden_states=a_		)




class 			__snake_case   ( SCREAMING_SNAKE_CASE):







							'''simple docstring'''






							UpperCamelCase__					:   int       =       ResNetConfig
							UpperCamelCase__					:   Tuple       =       """resnet"""
							UpperCamelCase__					:   Tuple       =       """pixel_values"""


							@property
							def _a   (				self		):
											return {"pixel_values": tf.TensorSpec(shape=(None, self.config.num_channels, 224, 224)	,	dtype=tf.floataa		)}




UpperCAmelCase							=       R"""
    This model is a TensorFlow
    [tf.keras.layers.Layer](https://www.tensorflow.org/api_docs/python/tf/keras/layers/Layer) sub-class. Use it as a
    regular TensorFlow Module and refer to the TensorFlow documentation for all matter related to general usage and
    behavior.

    Parameters:
        config ([`ResNetConfig`]): Model configuration class with all the parameters of the model.
            Initializing with a config file does not load the weights associated with the model, only the
            configuration. Check out the [`~TFPreTrainedModel.from_pretrained`] method to load the model weights.
"""


UpperCAmelCase							=       R"""
    Args:
        pixel_values (`tf.Tensor` of shape `(batch_size, num_channels, height, width)`):
            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See
            [`ConvNextImageProcessor.__call__`] for details.

        output_hidden_states (`bool`, *optional*):
            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
            more detail.
        return_dict (`bool`, *optional*):
            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
"""

@keras_serializable
class 			__snake_case   ( tf.keras.layers.Layer):







							'''simple docstring'''






							UpperCamelCase__					:   Union[str, Any]       =       ResNetConfig


							def __init__(				self	,	a_	,	**a_		):
											super().__init__(**a_		)
											a__					   =		config
											a__					   =		TFResNetEmbeddings(a_	,	name="""embedder"""		)
											a__					   =		TFResNetEncoder(a_	,	name="""encoder"""		)
											a__					   =		tf.keras.layers.GlobalAveragePoolingaD(keepdims=a_		)







							@unpack_inputs
							def _a   (				self	,	a_	,	a_ = None	,	a_ = None	,	a_ = False	,	):
											a__					   =		(
											    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
											)
											a__					   =		return_dict if return_dict is not None else self.config.use_return_dict

											# TF 2.0 image layers can't use NCHW format when running on CPU.
											# We transpose to NHWC format and then transpose back after the full forward pass.
											# (batch_size, num_channels, height, width) -> (batch_size, height, width, num_channels)
											a__					   =		tf.transpose(a_	,	perm=[0, 2, 3, 1]		)
											a__					   =		self.embedder(a_	,	training=a_		)

											a__					   =		self.encoder(
											    a_	,	output_hidden_states=a_	,	return_dict=a_	,	training=a_		)

											a__					   =		encoder_outputs[0]

											a__					   =		self.pooler(a_		)

											# Transpose all the outputs to the NCHW format
											# (batch_size, height, width, num_channels) -> (batch_size, num_channels, height, width)
											a__					   =		tf.transpose(a_	,	(0, 3, 1, 2)		)
											a__					   =		tf.transpose(a_	,	(0, 3, 1, 2)		)
											a__					   =		()
											for hidden_state in encoder_outputs[1:]:
															a__					   =		hidden_states + tuple(tf.transpose(a_	,	(0, 3, 1, 2)		) for h in hidden_state		)

											if not return_dict:
															return (last_hidden_state, pooled_output) + hidden_states

											a__					   =		hidden_states if output_hidden_states else None

											return TFBaseModelOutputWithPoolingAndNoAttention(
											    last_hidden_state=a_	,	pooler_output=a_	,	hidden_states=a_	,	)




@add_start_docstrings(
    """The bare ResNet model outputting raw features without any specific head on top."""      ,SCREAMING_SNAKE_CASE      ,)
class 			__snake_case   ( SCREAMING_SNAKE_CASE):







							'''simple docstring'''






							def __init__(				self	,	a_	,	**a_		):
											super().__init__(a_	,	**a_		)
											a__					   =		TFResNetMainLayer(config=a_	,	name="""resnet"""		)







							@add_start_docstrings_to_model_forward(a_		)
							@add_code_sample_docstrings(
							    checkpoint=_CHECKPOINT_FOR_DOC	,	output_type=a_	,	config_class=_CONFIG_FOR_DOC	,	modality="""vision"""	,	expected_output=_EXPECTED_OUTPUT_SHAPE	,	)
							@unpack_inputs
							def _a   (				self	,	a_	,	a_ = None	,	a_ = None	,	a_ = False	,	):
											a__					   =		(
											    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
											)
											a__					   =		return_dict if return_dict is not None else self.config.use_return_dict

											a__					   =		self.resnet(
											    pixel_values=a_	,	output_hidden_states=a_	,	return_dict=a_	,	training=a_	,	)
											return resnet_outputs




@add_start_docstrings(
    """
    ResNet Model with an image classification head on top (a linear layer on top of the pooled features), e.g. for
    ImageNet.
    """      ,SCREAMING_SNAKE_CASE      ,)
class 			__snake_case   ( SCREAMING_SNAKE_CASE      ,SCREAMING_SNAKE_CASE):







							'''simple docstring'''






							def __init__(				self	,	a_	,	**a_		):
											super().__init__(a_	,	**a_		)
											a__					   =		config.num_labels
											a__					   =		TFResNetMainLayer(a_	,	name="""resnet"""		)
											# classification head
											a__					   =		(
											    tf.keras.layers.Dense(config.num_labels	,	name="""classifier.1"""		)
											    if config.num_labels > 0
											    else tf.keras.layers.Activation("""linear"""	,	name="""classifier.1"""		)
											)


							def _a   (				self	,	a_		):
											a__					   =		tf.keras.layers.Flatten()(a_		)
											a__					   =		self.classifier_layer(a_		)
											return logits







							@add_start_docstrings_to_model_forward(a_		)
							@add_code_sample_docstrings(
							    checkpoint=_IMAGE_CLASS_CHECKPOINT	,	output_type=a_	,	config_class=_CONFIG_FOR_DOC	,	expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT	,	)
							@unpack_inputs
							def _a   (				self	,	a_ = None	,	a_ = None	,	a_ = None	,	a_ = None	,	a_ = False	,	):
											a__					   =		return_dict if return_dict is not None else self.config.use_return_dict

											a__					   =		self.resnet(
											    a_	,	output_hidden_states=a_	,	return_dict=a_	,	training=a_		)

											a__					   =		outputs.pooler_output if return_dict else outputs[1]

											a__					   =		self.classifier(a_		)

											a__					   =		None if labels is None else self.hf_compute_loss(a_	,	a_		)

											if not return_dict:
															a__					   =		(logits,) + outputs[2:]
															return (loss,) + output if loss is not None else output

											return TFImageClassifierOutputWithNoAttention(loss=a_	,	logits=a_	,	hidden_states=outputs.hidden_states		)








"""simple docstring"""




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

import numpy as np
import PIL
import torch
from transformers import CLIPTextModelWithProjection, CLIPTokenizer

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


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

UpperCamelCase_	 ="""
    Examples:
        ```py
        >>> import torch
        >>> from diffusers import DiffusionPipeline
        >>> from diffusers.utils import export_to_gif

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

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

        >>> guidance_scale = 15.0
        >>> prompt = \"a shark\"

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

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




@dataclass
class 					_a		( _lowerCAmelCase							):
				UpperCamelCase								=       42




class 					_a		( _lowerCAmelCase							):






				def __init__(					self						:       int,					lowerCAmelCase__						:       PriorTransformer,					lowerCAmelCase__						:       CLIPTextModelWithProjection,					lowerCAmelCase__						:       CLIPTokenizer,					lowerCAmelCase__						:       HeunDiscreteScheduler,					lowerCAmelCase__						:       ShapERenderer,					)      -> Optional[int]:




										'''simple docstring'''



										super().__init__()

										self.register_modules(
										    prior=lowerCAmelCase__,					text_encoder=lowerCAmelCase__,					tokenizer=lowerCAmelCase__,					scheduler=lowerCAmelCase__,					renderer=lowerCAmelCase__,					)






				def 							snake_case      (					self						:       List[Any],					lowerCAmelCase__						:       Any,					lowerCAmelCase__						:       Optional[Any],					lowerCAmelCase__						:       Optional[int],					lowerCAmelCase__						:       List[str],					lowerCAmelCase__						:       int,					lowerCAmelCase__						:       List[str]				)      -> Optional[int]:




										'''simple docstring'''



										if latents is None:
																_UpperCamelCase     :   Any      =       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}"""				)
																_UpperCamelCase     :   List[Any]      =       latents.to(lowerCAmelCase__				)

										_UpperCamelCase     :   Union[str, Any]      =       latents * scheduler.init_noise_sigma
										return latents






				def 							snake_case      (					self						:       int,					lowerCAmelCase__						:       str=0				)      -> Dict:




										'''simple docstring'''



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

										_UpperCamelCase     :   List[str]      =       torch.device(f"""cuda:{gpu_id}"""				)

										_UpperCamelCase     :   List[str]      =       [self.text_encoder, self.prior]
										for cpu_offloaded_model in models:
																if cpu_offloaded_model is not None:
																						cpu_offload(lowerCAmelCase__,					lowerCAmelCase__				)






				def 							snake_case      (					self						:       int,					lowerCAmelCase__						:       Tuple=0				)      -> List[Any]:




										'''simple docstring'''



										if is_accelerate_available() and is_accelerate_version('''>=''',					'''0.17.0.dev0'''				):
																from accelerate import cpu_offload_with_hook
										else:
																raise ImportError('''`enable_model_cpu_offload` requires `accelerate v0.17.0` or higher.'''				)

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

										if self.device.type != "cpu":
																self.to('''cpu''',					silence_dtype_warnings=lowerCAmelCase__				)
																torch.cuda.empty_cache()  # otherwise we don't see the memory savings (but they probably exist)

										_UpperCamelCase     :   Tuple      =       None
										for cpu_offloaded_model in [self.text_encoder, self.prior, self.renderer]:
																_UpperCamelCase  ,				_UpperCamelCase     :   List[Any]      =       cpu_offload_with_hook(lowerCAmelCase__,					lowerCAmelCase__,					prev_module_hook=lowerCAmelCase__				)

										if self.safety_checker is not None:
																_UpperCamelCase  ,				_UpperCamelCase     :   List[Any]      =       cpu_offload_with_hook(self.safety_checker,					lowerCAmelCase__,					prev_module_hook=lowerCAmelCase__				)

										# We'll offload the last model manually.
										_UpperCamelCase     :   Union[str, Any]      =       hook






				@property
				def 							snake_case      (					self						:       Any				)      -> Tuple:




										'''simple docstring'''



										if self.device != torch.device('''meta'''				) or not hasattr(self.text_encoder,					'''_hf_hook'''				):
																return self.device
										for module in self.text_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 							snake_case      (					self						:       Tuple,					lowerCAmelCase__						:       Optional[Any],					lowerCAmelCase__						:       int,					lowerCAmelCase__						:       Tuple,					lowerCAmelCase__						:       Any,					)      -> List[str]:




										'''simple docstring'''



										len(lowerCAmelCase__				) if isinstance(lowerCAmelCase__,					lowerCAmelCase__				) else 1

										# YiYi Notes: set pad_token_id to be 0, not sure why I can't set in the config file
										_UpperCamelCase     :   Optional[Any]      =       0
										# get prompt text embeddings
										_UpperCamelCase     :   Union[str, Any]      =       self.tokenizer(
										    lowerCAmelCase__,					padding='''max_length''',					max_length=self.tokenizer.model_max_length,					truncation=lowerCAmelCase__,					return_tensors='''pt''',					)
										_UpperCamelCase     :   Union[str, Any]      =       text_inputs.input_ids
										_UpperCamelCase     :   Union[str, Any]      =       self.tokenizer(lowerCAmelCase__,					padding='''longest''',					return_tensors='''pt'''				).input_ids

										if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(lowerCAmelCase__,					lowerCAmelCase__				):
																_UpperCamelCase     :   Tuple      =       self.tokenizer.batch_decode(untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1]				)
																logger.warning(
																    '''The following part of your input was truncated because CLIP can only handle sequences up to'''
																    f""" {self.tokenizer.model_max_length} tokens: {removed_text}"""				)

										_UpperCamelCase     :   Dict      =       self.text_encoder(text_input_ids.to(lowerCAmelCase__				)				)
										_UpperCamelCase     :   Optional[int]      =       text_encoder_output.text_embeds

										_UpperCamelCase     :   Any      =       prompt_embeds.repeat_interleave(lowerCAmelCase__,					dim=0				)
										# in Shap-E it normalize the prompt_embeds and then later rescale it
										_UpperCamelCase     :   str      =       prompt_embeds / torch.linalg.norm(lowerCAmelCase__,					dim=-1,					keepdim=lowerCAmelCase__				)

										if do_classifier_free_guidance:
																_UpperCamelCase     :   Any      =       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
																_UpperCamelCase     :   Optional[Any]      =       torch.cat([negative_prompt_embeds, prompt_embeds]				)

										# Rescale the features to have unit variance
										_UpperCamelCase     :   Union[str, Any]      =       math.sqrt(prompt_embeds.shape[1]				) * prompt_embeds

										return prompt_embeds






				@torch.no_grad()
				@replace_example_docstring(lowerCAmelCase__				)
				def __call__(					self						:       Optional[int],					lowerCAmelCase__						:       str,					lowerCAmelCase__						:       int = 1,					lowerCAmelCase__						:       int = 2_5,					lowerCAmelCase__						:       Optional[Union[torch.Generator, List[torch.Generator]]] = None,					lowerCAmelCase__						:       Optional[torch.FloatTensor] = None,					lowerCAmelCase__						:       float = 4.0,					lowerCAmelCase__						:       int = 6_4,					lowerCAmelCase__						:       Optional[str] = "pil",					lowerCAmelCase__						:       bool = True,					)      -> Any:




										'''simple docstring'''



										if isinstance(lowerCAmelCase__,					lowerCAmelCase__				):
																_UpperCamelCase     :   Any      =       1
										elif isinstance(lowerCAmelCase__,					lowerCAmelCase__				):
																_UpperCamelCase     :   str      =       len(lowerCAmelCase__				)
										else:
																raise ValueError(f"""`prompt` has to be of type `str` or `list` but is {type(lowerCAmelCase__				)}"""				)

										_UpperCamelCase     :   Dict      =       self._execution_device

										_UpperCamelCase     :   List[Any]      =       batch_size * num_images_per_prompt

										_UpperCamelCase     :   List[str]      =       guidance_scale > 1.0
										_UpperCamelCase     :   str      =       self._encode_prompt(lowerCAmelCase__,					lowerCAmelCase__,					lowerCAmelCase__,					lowerCAmelCase__				)

										# prior

										self.scheduler.set_timesteps(lowerCAmelCase__,					device=lowerCAmelCase__				)
										_UpperCamelCase     :   Any      =       self.scheduler.timesteps

										_UpperCamelCase     :   Optional[Any]      =       self.prior.config.num_embeddings
										_UpperCamelCase     :   Optional[Any]      =       self.prior.config.embedding_dim

										_UpperCamelCase     :   Tuple      =       self.prepare_latents(
										    (batch_size, num_embeddings * embedding_dim),					prompt_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
										_UpperCamelCase     :   Dict      =       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
																_UpperCamelCase     :   Optional[int]      =       torch.cat([latents] * 2				) if do_classifier_free_guidance else latents
																_UpperCamelCase     :   Dict      =       self.scheduler.scale_model_input(lowerCAmelCase__,					lowerCAmelCase__				)

																_UpperCamelCase     :   Tuple      =       self.prior(
																    lowerCAmelCase__,					timestep=lowerCAmelCase__,					proj_embedding=lowerCAmelCase__,					).predicted_image_embedding

																# remove the variance
																_UpperCamelCase  ,				_UpperCamelCase     :   Tuple      =       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     :   List[str]      =       noise_pred.chunk(2				)
																						_UpperCamelCase     :   Optional[Any]      =       noise_pred_uncond + guidance_scale * (noise_pred - noise_pred_uncond)

																_UpperCamelCase     :   Dict      =       self.scheduler.step(
																    lowerCAmelCase__,					timestep=lowerCAmelCase__,					sample=lowerCAmelCase__,					).prev_sample

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

										_UpperCamelCase     :   List[Any]      =       []
										for i, latent in enumerate(lowerCAmelCase__				):
																_UpperCamelCase     :   Union[str, Any]      =       self.renderer.decode(
																    latent[None, :],					lowerCAmelCase__,					size=lowerCAmelCase__,					ray_batch_size=4_0_9_6,					n_coarse_samples=6_4,					n_fine_samples=1_2_8,					)
																images.append(lowerCAmelCase__				)

										_UpperCamelCase     :   str      =       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}"""				)

										_UpperCamelCase     :   Dict      =       images.cpu().numpy()

										if output_type == "pil":
																_UpperCamelCase     :   Any      =       [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__				)



