code: 
  
  
  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
  
  
  SCREAMING_SNAKE_CASE__				:    List[Any]         =    logging.get_logger(__name__)
  
  # General docstring
  SCREAMING_SNAKE_CASE__				:    Optional[int]         =    'ResNetConfig'
  
  # Base docstring
  SCREAMING_SNAKE_CASE__				:    Any         =    'microsoft/resnet-50'
  SCREAMING_SNAKE_CASE__				:    str         =    [1, 2_048, 7, 7]
  
  # Image classification docstring
  SCREAMING_SNAKE_CASE__				:    Any         =    'microsoft/resnet-50'
  SCREAMING_SNAKE_CASE__				:    str         =    'tiger cat'
  
  SCREAMING_SNAKE_CASE__				:    List[Any]         =    [
      'microsoft/resnet-50',
      # See all resnet models at https://huggingface.co/models?filter=resnet
  ]
  
  
  class 		_SCREAMING_SNAKE_CASE    (	tf.keras.layers.Layer    ):
  
  
       def __init__(     self       ,				A_       ,				A_ = 3       ,				A_ = 1       ,				A_ = "relu"       ,				**A_						):
         super().__init__(**A_						)
         _UpperCAmelCase		:       Union[str, Any]       =    kernel_size // 2
         _UpperCAmelCase		:       Tuple       =    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
         _UpperCAmelCase		:       Optional[int]       =    tf.keras.layers.BatchNormalization(epsilon=1e-5       ,				momentum=0.9       ,				name="""normalization"""						)
         _UpperCAmelCase		:       Tuple       =    ACTaFN[activation] if activation is not None else tf.keras.layers.Activation("""linear"""						)
  
  
       def 					__snake_case(     self       ,				A_						):
         # Pad to match that done in the PyTorch Conv2D model
         _UpperCAmelCase		:       Any       =    (self.pad_value, self.pad_value)
         _UpperCAmelCase		:       Tuple       =    tf.pad(A_       ,				[(0, 0), height_pad, width_pad, (0, 0)]						)
         _UpperCAmelCase		:       List[Any]       =    self.conv(A_						)
         return hidden_state
  
  
  
  
       def 					__snake_case(     self       ,				A_       ,				A_ = False						):
         _UpperCAmelCase		:       Any       =    self.convolution(A_						)
         _UpperCAmelCase		:       Optional[Any]       =    self.normalization(A_       ,				training=A_						)
         _UpperCAmelCase		:       List[str]       =    self.activation(A_						)
         return hidden_state
  
  
  class 		_SCREAMING_SNAKE_CASE    (	tf.keras.layers.Layer    ):
  
  
       def __init__(     self       ,				A_       ,				**A_						):
         super().__init__(**A_						)
         _UpperCAmelCase		:       str       =    TFResNetConvLayer(
             config.embedding_size       ,				kernel_size=7       ,				stride=2       ,				activation=config.hidden_act       ,				name="""embedder"""       ,				)
         _UpperCAmelCase		:       Any       =    tf.keras.layers.MaxPoolaD(pool_size=3       ,				strides=2       ,				padding="""valid"""       ,				name="""pooler"""						)
         _UpperCAmelCase		:       Dict       =    config.num_channels
  
  
  
  
       def 					__snake_case(     self       ,				A_       ,				A_ = False						):
         _UpperCAmelCase		:       Optional[int]       =    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."""						)
         _UpperCAmelCase		:       Optional[Any]       =    pixel_values
         _UpperCAmelCase		:       List[str]       =    self.embedder(A_						)
         _UpperCAmelCase		:       str       =    tf.pad(A_       ,				[[0, 0], [1, 1], [1, 1], [0, 0]]						)
         _UpperCAmelCase		:       List[Any]       =    self.pooler(A_						)
         return hidden_state
  
  
  class 		_SCREAMING_SNAKE_CASE    (	tf.keras.layers.Layer    ):
  
  
       def __init__(     self       ,				A_       ,				A_ = 2       ,				**A_						):
         super().__init__(**A_						)
         _UpperCAmelCase		:       Optional[int]       =    tf.keras.layers.ConvaD(
             A_       ,				kernel_size=1       ,				strides=A_       ,				use_bias=A_       ,				name="""convolution"""						)
         # Use same default momentum and epsilon as PyTorch equivalent
         _UpperCAmelCase		:       List[str]       =    tf.keras.layers.BatchNormalization(epsilon=1e-5       ,				momentum=0.9       ,				name="""normalization"""						)
  
  
  
  
       def 					__snake_case(     self       ,				A_       ,				A_ = False						):
         _UpperCAmelCase		:       Union[str, Any]       =    x
         _UpperCAmelCase		:       Optional[int]       =    self.convolution(A_						)
         _UpperCAmelCase		:       int       =    self.normalization(A_       ,				training=A_						)
         return hidden_state
  
  
  class 		_SCREAMING_SNAKE_CASE    (	tf.keras.layers.Layer    ):
  
  
       def __init__(     self       ,				A_       ,				A_       ,				A_ = 1       ,				A_ = "relu"       ,				**A_						):
         super().__init__(**A_						)
         _UpperCAmelCase		:       Optional[int]       =    in_channels != out_channels or stride != 1
         _UpperCAmelCase		:       str       =    TFResNetConvLayer(A_       ,				stride=A_       ,				name="""layer.0"""						)
         _UpperCAmelCase		:       Optional[Any]       =    TFResNetConvLayer(A_       ,				activation=A_       ,				name="""layer.1"""						)
         _UpperCAmelCase		:       Union[str, Any]       =    (
             TFResNetShortCut(A_       ,				stride=A_       ,				name="""shortcut"""						)
             if should_apply_shortcut
             else tf.keras.layers.Activation("""linear"""       ,				name="""shortcut"""						)
         )
         _UpperCAmelCase		:       Tuple       =    ACTaFN[activation]
  
  
  
  
       def 					__snake_case(     self       ,				A_       ,				A_ = False						):
         _UpperCAmelCase		:       Dict       =    hidden_state
         _UpperCAmelCase		:       Union[str, Any]       =    self.conva(A_       ,				training=A_						)
         _UpperCAmelCase		:       Union[str, Any]       =    self.conva(A_       ,				training=A_						)
         _UpperCAmelCase		:       int       =    self.shortcut(A_       ,				training=A_						)
         hidden_state += residual
         _UpperCAmelCase		:       List[str]       =    self.activation(A_						)
         return hidden_state
  
  
  class 		_SCREAMING_SNAKE_CASE    (	tf.keras.layers.Layer    ):
  
  
       def __init__(     self       ,				A_       ,				A_       ,				A_ = 1       ,				A_ = "relu"       ,				A_ = 4       ,				**A_       ,				):
         super().__init__(**A_						)
         _UpperCAmelCase		:       Any       =    in_channels != out_channels or stride != 1
         _UpperCAmelCase		:       List[str]       =    out_channels // reduction
         _UpperCAmelCase		:       Any       =    TFResNetConvLayer(A_       ,				kernel_size=1       ,				name="""layer.0"""						)
         _UpperCAmelCase		:       List[Any]       =    TFResNetConvLayer(A_       ,				stride=A_       ,				name="""layer.1"""						)
         _UpperCAmelCase		:       Tuple       =    TFResNetConvLayer(A_       ,				kernel_size=1       ,				activation=A_       ,				name="""layer.2"""						)
         _UpperCAmelCase		:       int       =    (
             TFResNetShortCut(A_       ,				stride=A_       ,				name="""shortcut"""						)
             if should_apply_shortcut
             else tf.keras.layers.Activation("""linear"""       ,				name="""shortcut"""						)
         )
         _UpperCAmelCase		:       int       =    ACTaFN[activation]
  
  
  
  
       def 					__snake_case(     self       ,				A_       ,				A_ = False						):
         _UpperCAmelCase		:       Tuple       =    hidden_state
         _UpperCAmelCase		:       str       =    self.conva(A_       ,				training=A_						)
         _UpperCAmelCase		:       List[Any]       =    self.conva(A_       ,				training=A_						)
         _UpperCAmelCase		:       int       =    self.conva(A_       ,				training=A_						)
         _UpperCAmelCase		:       List[Any]       =    self.shortcut(A_       ,				training=A_						)
         hidden_state += residual
         _UpperCAmelCase		:       str       =    self.activation(A_						)
         return hidden_state
  
  
  class 		_SCREAMING_SNAKE_CASE    (	tf.keras.layers.Layer    ):
  
  
       def __init__(     self       ,				A_       ,				A_       ,				A_       ,				A_ = 2       ,				A_ = 2       ,				**A_						):
         super().__init__(**A_						)
  
         _UpperCAmelCase		:       Optional[int]       =    TFResNetBottleNeckLayer if config.layer_type == """bottleneck""" else TFResNetBasicLayer
  
         _UpperCAmelCase		:       Tuple       =    [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						)
         ]
         _UpperCAmelCase		:       Any       =    layers
  
  
  
  
       def 					__snake_case(     self       ,				A_       ,				A_ = False						):
         for layer in self.stage_layers:
           _UpperCAmelCase		:       Union[str, Any]       =    layer(A_       ,				training=A_						)
         return hidden_state
  
  
  class 		_SCREAMING_SNAKE_CASE    (	tf.keras.layers.Layer    ):
  
  
       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
         _UpperCAmelCase		:       Dict       =    [
             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 					__snake_case(     self       ,				A_       ,				A_ = False       ,				A_ = True       ,				A_ = False       ,				):
         _UpperCAmelCase		:       Dict       =    () if output_hidden_states else None
  
         for stage_module in self.stages:
           if output_hidden_states:
             _UpperCAmelCase		:       int       =    hidden_states + (hidden_state,)
  
           _UpperCAmelCase		:       List[Any]       =    stage_module(A_       ,				training=A_						)
  
         if output_hidden_states:
           _UpperCAmelCase		:       Dict       =    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 		_SCREAMING_SNAKE_CASE    (	A    ):
       __SCREAMING_SNAKE_CASE    				=						ResNetConfig
       __SCREAMING_SNAKE_CASE    				=						'''resnet'''
       __SCREAMING_SNAKE_CASE    				=						'''pixel_values'''
  
  
       @property
       def 					__snake_case(     self						):
         return {"pixel_values": tf.TensorSpec(shape=(None, self.config.num_channels, 2_24, 2_24)       ,				dtype=tf.floataa						)}
  
  
  SCREAMING_SNAKE_CASE__				:    List[Any]         =    R'\n    This model is a TensorFlow\n    [tf.keras.layers.Layer](https://www.tensorflow.org/api_docs/python/tf/keras/layers/Layer) sub-class. Use it as a\n    regular TensorFlow Module and refer to the TensorFlow documentation for all matter related to general usage and\n    behavior.\n\n    Parameters:\n        config ([`ResNetConfig`]): Model configuration class with all the parameters of the model.\n            Initializing with a config file does not load the weights associated with the model, only the\n            configuration. Check out the [`~TFPreTrainedModel.from_pretrained`] method to load the model weights.\n'
  
  
  SCREAMING_SNAKE_CASE__				:    str         =    R'\n    Args:\n        pixel_values (`tf.Tensor` of shape `(batch_size, num_channels, height, width)`):\n            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See\n            [`ConvNextImageProcessor.__call__`] for details.\n\n        output_hidden_states (`bool`, *optional*):\n            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for\n            more detail.\n        return_dict (`bool`, *optional*):\n            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.\n'
  
  
  @keras_serializable
  class 		_SCREAMING_SNAKE_CASE    (	tf.keras.layers.Layer    ):
       __SCREAMING_SNAKE_CASE    				=						ResNetConfig
  
  
       def __init__(     self       ,				A_       ,				**A_						):
         super().__init__(**A_						)
         _UpperCAmelCase		:       int       =    config
         _UpperCAmelCase		:       str       =    TFResNetEmbeddings(A_       ,				name="""embedder"""						)
         _UpperCAmelCase		:       int       =    TFResNetEncoder(A_       ,				name="""encoder"""						)
         _UpperCAmelCase		:       List[Any]       =    tf.keras.layers.GlobalAveragePoolingaD(keepdims=A_						)
  
  
  
  
       @unpack_inputs
       def 					__snake_case(     self       ,				A_       ,				A_ = None       ,				A_ = None       ,				A_ = False       ,				):
         _UpperCAmelCase		:       Optional[Any]       =    (
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
         )
         _UpperCAmelCase		:       int       =    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)
         _UpperCAmelCase		:       Any       =    tf.transpose(A_       ,				perm=[0, 2, 3, 1]						)
         _UpperCAmelCase		:       Union[str, Any]       =    self.embedder(A_       ,				training=A_						)
  
         _UpperCAmelCase		:       Tuple       =    self.encoder(
             A_       ,				output_hidden_states=A_       ,				return_dict=A_       ,				training=A_						)
  
         _UpperCAmelCase		:       Any       =    encoder_outputs[0]
  
         _UpperCAmelCase		:       Any       =    self.pooler(A_						)
  
         # Transpose all the outputs to the NCHW format
         # (batch_size, height, width, num_channels) -> (batch_size, num_channels, height, width)
         _UpperCAmelCase		:       Any       =    tf.transpose(A_       ,				(0, 3, 1, 2)						)
         _UpperCAmelCase		:       Dict       =    tf.transpose(A_       ,				(0, 3, 1, 2)						)
         _UpperCAmelCase		:       int       =    ()
         for hidden_state in encoder_outputs[1:]:
           _UpperCAmelCase		:       Tuple       =    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
  
         _UpperCAmelCase		:       Tuple       =    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.'''       ,  A       ,  )
  class 		_SCREAMING_SNAKE_CASE    (	A    ):
  
  
       def __init__(     self       ,				A_       ,				**A_						):
         super().__init__(A_       ,				**A_						)
         _UpperCAmelCase		:       int       =    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 					__snake_case(     self       ,				A_       ,				A_ = None       ,				A_ = None       ,				A_ = False       ,				):
         _UpperCAmelCase		:       int       =    (
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
         )
         _UpperCAmelCase		:       Union[str, Any]       =    return_dict if return_dict is not None else self.config.use_return_dict
  
         _UpperCAmelCase		:       List[str]       =    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.
      '''       ,  A       ,  )
  class 		_SCREAMING_SNAKE_CASE    (	A       ,  A    ):
  
  
       def __init__(     self       ,				A_       ,				**A_						):
         super().__init__(A_       ,				**A_						)
         _UpperCAmelCase		:       int       =    config.num_labels
         _UpperCAmelCase		:       str       =    TFResNetMainLayer(A_       ,				name="""resnet"""						)
         # classification head
         _UpperCAmelCase		:       str       =    (
             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 					__snake_case(     self       ,				A_						):
         _UpperCAmelCase		:       int       =    tf.keras.layers.Flatten()(A_						)
         _UpperCAmelCase		:       str       =    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 					__snake_case(     self       ,				A_ = None       ,				A_ = None       ,				A_ = None       ,				A_ = None       ,				A_ = False       ,				):
         _UpperCAmelCase		:       List[Any]       =    return_dict if return_dict is not None else self.config.use_return_dict
  
         _UpperCAmelCase		:       List[Any]       =    self.resnet(
             A_       ,				output_hidden_states=A_       ,				return_dict=A_       ,				training=A_						)
  
         _UpperCAmelCase		:       Optional[int]       =    outputs.pooler_output if return_dict else outputs[1]
  
         _UpperCAmelCase		:       str       =    self.classifier(A_						)
  
         _UpperCAmelCase		:       Optional[Any]       =    None if labels is None else self.hf_compute_loss(A_       ,				A_						)
  
         if not return_dict:
           _UpperCAmelCase		:       Optional[int]       =    (logits,) + outputs[2:]
           return (loss,) + output if loss is not None else output
  
         return TFImageClassifierOutputWithNoAttention(loss=A_       ,				logits=A_       ,				hidden_states=outputs.hidden_states						)
  
  
  
  
  
  
style_context: 
  
  
  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
  
  
  SCREAMING_SNAKE_CASE__				:    Optional[int]         =    logging.get_logger(__name__)  # pylint: disable=invalid-name
  
  SCREAMING_SNAKE_CASE__				:    Tuple         =    '\n    Examples:\n        ```py\n        >>> import torch\n        >>> from diffusers import DiffusionPipeline\n        >>> from diffusers.utils import export_to_gif\n\n        >>> device = torch.device("cuda" if torch.cuda.is_available() else "cpu")\n\n        >>> repo = "openai/shap-e"\n        >>> pipe = DiffusionPipeline.from_pretrained(repo, torch_dtype=torch.float16)\n        >>> pipe = pipe.to(device)\n\n        >>> guidance_scale = 15.0\n        >>> prompt = "a shark"\n\n        >>> images = pipe(\n        ...     prompt,\n        ...     guidance_scale=guidance_scale,\n        ...     num_inference_steps=64,\n        ...     frame_size=256,\n        ... ).images\n\n        >>> gif_path = export_to_gif(images[0], "shark_3d.gif")\n        ```\n'
  
  
  @dataclass
  class 		_SCREAMING_SNAKE_CASE    (	A    ):
       __SCREAMING_SNAKE_CASE    				=						42
  
  
  class 		_SCREAMING_SNAKE_CASE    (	A    ):
  
  
       def __init__(     self       ,				A_       ,				A_       ,				A_       ,				A_       ,				A_       ,				):
         super().__init__()
  
         self.register_modules(
             prior=A_       ,				text_encoder=A_       ,				tokenizer=A_       ,				scheduler=A_       ,				renderer=A_       ,				)
  
  
       def 					__snake_case(     self       ,				A_       ,				A_       ,				A_       ,				A_       ,				A_       ,				A_						):
         if latents is None:
           _UpperCAmelCase		:       List[Any]       =    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		:       Optional[int]       =    latents.to(A_						)
  
         _UpperCAmelCase		:       str       =    latents * scheduler.init_noise_sigma
         return latents
  
  
       def 					__snake_case(     self       ,				A_=0						):
         if is_accelerate_available():
           from accelerate import cpu_offload
         else:
           raise ImportError("""Please install accelerate via `pip install accelerate`"""						)
  
         _UpperCAmelCase		:       Any       =    torch.device(F'''cuda:{gpu_id}'''						)
  
         _UpperCAmelCase		:       Dict       =    [self.text_encoder, self.prior]
         for cpu_offloaded_model in models:
           if cpu_offloaded_model is not None:
             cpu_offload(A_       ,				A_						)
  
  
       def 					__snake_case(     self       ,				A_=0						):
         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		:       int       =    torch.device(F'''cuda:{gpu_id}'''						)
  
         if self.device.type != "cpu":
           self.to("""cpu"""       ,				silence_dtype_warnings=A_						)
           torch.cuda.empty_cache()  # otherwise we don't see the memory savings (but they probably exist)
  
         _UpperCAmelCase		:       List[str]       =    None
         for cpu_offloaded_model in [self.text_encoder, self.prior, self.renderer]:
           _UpperCAmelCase,_UpperCAmelCase		:       List[Any]       =    cpu_offload_with_hook(A_       ,				A_       ,				prev_module_hook=A_						)
  
         if self.safety_checker is not None:
           _UpperCAmelCase,_UpperCAmelCase		:       int       =    cpu_offload_with_hook(self.safety_checker       ,				A_       ,				prev_module_hook=A_						)
  
         # We'll offload the last model manually.
         _UpperCAmelCase		:       Optional[Any]       =    hook
  
  
       @property
       def 					__snake_case(     self						):
         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(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 					__snake_case(     self       ,				A_       ,				A_       ,				A_       ,				A_       ,				):
         len(A_						) if isinstance(A_       ,				A_						) else 1
  
         # YiYi Notes: set pad_token_id to be 0, not sure why I can't set in the config file
         _UpperCAmelCase		:       str       =    0
         # get prompt text embeddings
         _UpperCAmelCase		:       List[str]       =    self.tokenizer(
             A_       ,				padding="""max_length"""       ,				max_length=self.tokenizer.model_max_length       ,				truncation=A_       ,				return_tensors="""pt"""       ,				)
         _UpperCAmelCase		:       Any       =    text_inputs.input_ids
         _UpperCAmelCase		:       Dict       =    self.tokenizer(A_       ,				padding="""longest"""       ,				return_tensors="""pt"""						).input_ids
  
         if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(A_       ,				A_						):
           _UpperCAmelCase		:       Union[str, Any]       =    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		:       List[Any]       =    self.text_encoder(text_input_ids.to(A_						)						)
         _UpperCAmelCase		:       Union[str, Any]       =    text_encoder_output.text_embeds
  
         _UpperCAmelCase		:       Tuple       =    prompt_embeds.repeat_interleave(A_       ,				dim=0						)
         # in Shap-E it normalize the prompt_embeds and then later rescale it
         _UpperCAmelCase		:       Any       =    prompt_embeds / torch.linalg.norm(A_       ,				dim=-1       ,				keepdim=A_						)
  
         if do_classifier_free_guidance:
           _UpperCAmelCase		:       Union[str, Any]       =    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		:       Optional[Any]       =    torch.cat([negative_prompt_embeds, prompt_embeds]						)
  
         # Rescale the features to have unit variance
         _UpperCAmelCase		:       List[Any]       =    math.sqrt(prompt_embeds.shape[1]						) * prompt_embeds
  
         return prompt_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       ,				):
  
         if isinstance(A_       ,				A_						):
           _UpperCAmelCase		:       Optional[int]       =    1
         elif isinstance(A_       ,				A_						):
           _UpperCAmelCase		:       Union[str, Any]       =    len(A_						)
         else:
           raise ValueError(F'''`prompt` has to be of type `str` or `list` but is {type(A_						)}'''						)
  
         _UpperCAmelCase		:       Tuple       =    self._execution_device
  
         _UpperCAmelCase		:       List[str]       =    batch_size * num_images_per_prompt
  
         _UpperCAmelCase		:       str       =    guidance_scale > 1.0
         _UpperCAmelCase		:       Tuple       =    self._encode_prompt(A_       ,				A_       ,				A_       ,				A_						)
  
         # prior
  
         self.scheduler.set_timesteps(A_       ,				device=A_						)
         _UpperCAmelCase		:       Any       =    self.scheduler.timesteps
  
         _UpperCAmelCase		:       List[str]       =    self.prior.config.num_embeddings
         _UpperCAmelCase		:       Union[str, Any]       =    self.prior.config.embedding_dim
  
         _UpperCAmelCase		:       Optional[Any]       =    self.prepare_latents(
             (batch_size, num_embeddings * embedding_dim)       ,				prompt_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		:       Optional[Any]       =    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		:       Optional[int]       =    torch.cat([latents] * 2						) if do_classifier_free_guidance else latents
           _UpperCAmelCase		:       Union[str, Any]       =    self.scheduler.scale_model_input(A_       ,				A_						)
  
           _UpperCAmelCase		:       Any       =    self.prior(
               A_       ,				timestep=A_       ,				proj_embedding=A_       ,				).predicted_image_embedding
  
           # remove the variance
           _UpperCAmelCase,_UpperCAmelCase		:       Dict       =    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		:       Union[str, Any]       =    noise_pred.chunk(2						)
             _UpperCAmelCase		:       List[Any]       =    noise_pred_uncond + guidance_scale * (noise_pred - noise_pred_uncond)
  
           _UpperCAmelCase		:       Tuple       =    self.scheduler.step(
               A_       ,				timestep=A_       ,				sample=A_       ,				).prev_sample
  
         if output_type == "latent":
           return ShapEPipelineOutput(images=A_						)
  
         _UpperCAmelCase		:       Dict       =    []
         for i, latent in enumerate(A_						):
           _UpperCAmelCase		:       int       =    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		:       List[Any]       =    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		:       str       =    images.cpu().numpy()
  
         if output_type == "pil":
           _UpperCAmelCase		:       List[str]       =    [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_						)
  
  
  
  
  
  