


'''simple docstring'''



import argparse
import os.path as osp
import re

import torch
from safetensors.torch import load_file, save_file


# =================#
# UNet Conversion #
# =================#

__UpperCAmelCase      :int			    =						[
    # (stable-diffusion, HF Diffusers)
    ("time_embed.0.weight", "time_embedding.linear_1.weight"),
    ("time_embed.0.bias", "time_embedding.linear_1.bias"),
    ("time_embed.2.weight", "time_embedding.linear_2.weight"),
    ("time_embed.2.bias", "time_embedding.linear_2.bias"),
    ("input_blocks.0.0.weight", "conv_in.weight"),
    ("input_blocks.0.0.bias", "conv_in.bias"),
    ("out.0.weight", "conv_norm_out.weight"),
    ("out.0.bias", "conv_norm_out.bias"),
    ("out.2.weight", "conv_out.weight"),
    ("out.2.bias", "conv_out.bias"),
]

__UpperCAmelCase      :List[Any]			    =						[
    # (stable-diffusion, HF Diffusers)
    ("in_layers.0", "norm1"),
    ("in_layers.2", "conv1"),
    ("out_layers.0", "norm2"),
    ("out_layers.3", "conv2"),
    ("emb_layers.1", "time_emb_proj"),
    ("skip_connection", "conv_shortcut"),
]

__UpperCAmelCase      :Dict			    =						[]
# hardcoded number of downblocks and resnets/attentions...
# would need smarter logic for other networks.
for i in range(4):
      # loop over downblocks/upblocks

      for j in range(2):
            # loop over resnets/attentions for downblocks
            __UpperCAmelCase      :str			    =						f"""down_blocks.{i}.resnets.{j}."""
            __UpperCAmelCase      :str			    =						f"""input_blocks.{3*i + j + 1}.0."""
            unet_conversion_map_layer.append((sd_down_res_prefix, hf_down_res_prefix))

            if i < 3:
                  # no attention layers in down_blocks.3
                  __UpperCAmelCase      :Union[str, Any]			    =						f"""down_blocks.{i}.attentions.{j}."""
                  __UpperCAmelCase      :Optional[Any]			    =						f"""input_blocks.{3*i + j + 1}.1."""
                  unet_conversion_map_layer.append((sd_down_atn_prefix, hf_down_atn_prefix))

      for j in range(3):
            # loop over resnets/attentions for upblocks
            __UpperCAmelCase      :List[Any]			    =						f"""up_blocks.{i}.resnets.{j}."""
            __UpperCAmelCase      :List[str]			    =						f"""output_blocks.{3*i + j}.0."""
            unet_conversion_map_layer.append((sd_up_res_prefix, hf_up_res_prefix))

            if i > 0:
                  # no attention layers in up_blocks.0
                  __UpperCAmelCase      :List[Any]			    =						f"""up_blocks.{i}.attentions.{j}."""
                  __UpperCAmelCase      :Union[str, Any]			    =						f"""output_blocks.{3*i + j}.1."""
                  unet_conversion_map_layer.append((sd_up_atn_prefix, hf_up_atn_prefix))

      if i < 3:
            # no downsample in down_blocks.3
            __UpperCAmelCase      :List[str]			    =						f"""down_blocks.{i}.downsamplers.0.conv."""
            __UpperCAmelCase      :Tuple			    =						f"""input_blocks.{3*(i+1)}.0.op."""
            unet_conversion_map_layer.append((sd_downsample_prefix, hf_downsample_prefix))

            # no upsample in up_blocks.3
            __UpperCAmelCase      :str			    =						f"""up_blocks.{i}.upsamplers.0."""
            __UpperCAmelCase      :Tuple			    =						f"""output_blocks.{3*i + 2}.{1 if i == 0 else 2}."""
            unet_conversion_map_layer.append((sd_upsample_prefix, hf_upsample_prefix))

__UpperCAmelCase      :Dict			    =						"mid_block.attentions.0."
__UpperCAmelCase      :Optional[Any]			    =						"middle_block.1."
unet_conversion_map_layer.append((sd_mid_atn_prefix, hf_mid_atn_prefix))

for j in range(2):
      __UpperCAmelCase      :Union[str, Any]			    =						f"""mid_block.resnets.{j}."""
      __UpperCAmelCase      :Optional[Any]			    =						f"""middle_block.{2*j}."""
      unet_conversion_map_layer.append((sd_mid_res_prefix, hf_mid_res_prefix))




def      _a     (						_lowercase :     Dict  ):





     '''simple docstring'''






     __UpperCAmelCase     :			Any  							=      {k: k for k in unet_state_dict.keys()}
     for sd_name, hf_name in unet_conversion_map:
          __UpperCAmelCase     :			Union[str, Any]  							=      sd_name
     for k, v in mapping.items():
          if "resnets" in k:
               for sd_part, hf_part in unet_conversion_map_resnet:
                    __UpperCAmelCase     :			Union[str, Any]  							=      v.replace(_lowercase      , _lowercase  )
               __UpperCAmelCase     :			Union[str, Any]  							=      v
     for k, v in mapping.items():
          for sd_part, hf_part in unet_conversion_map_layer:
               __UpperCAmelCase     :			List[str]  							=      v.replace(_lowercase      , _lowercase  )
          __UpperCAmelCase     :			List[str]  							=      v
     __UpperCAmelCase     :			Dict  							=      {v: unet_state_dict[k] for k, v in mapping.items()}
     return new_state_dict


# ================#
# VAE Conversion #
# ================#

__UpperCAmelCase      :List[str]			    =						[
    # (stable-diffusion, HF Diffusers)
    ("nin_shortcut", "conv_shortcut"),
    ("norm_out", "conv_norm_out"),
    ("mid.attn_1.", "mid_block.attentions.0."),
]

for i in range(4):
      # down_blocks have two resnets
      for j in range(2):
            __UpperCAmelCase      :Optional[int]			    =						f"""encoder.down_blocks.{i}.resnets.{j}."""
            __UpperCAmelCase      :Dict			    =						f"""encoder.down.{i}.block.{j}."""
            vae_conversion_map.append((sd_down_prefix, hf_down_prefix))

      if i < 3:
            __UpperCAmelCase      :List[Any]			    =						f"""down_blocks.{i}.downsamplers.0."""
            __UpperCAmelCase      :List[str]			    =						f"""down.{i}.downsample."""
            vae_conversion_map.append((sd_downsample_prefix, hf_downsample_prefix))

            __UpperCAmelCase      :int			    =						f"""up_blocks.{i}.upsamplers.0."""
            __UpperCAmelCase      :List[Any]			    =						f"""up.{3-i}.upsample."""
            vae_conversion_map.append((sd_upsample_prefix, hf_upsample_prefix))

      # up_blocks have three resnets
      # also, up blocks in hf are numbered in reverse from sd
      for j in range(3):
            __UpperCAmelCase      :str			    =						f"""decoder.up_blocks.{i}.resnets.{j}."""
            __UpperCAmelCase      :List[Any]			    =						f"""decoder.up.{3-i}.block.{j}."""
            vae_conversion_map.append((sd_up_prefix, hf_up_prefix))

# this part accounts for mid blocks in both the encoder and the decoder
for i in range(2):
      __UpperCAmelCase      :str			    =						f"""mid_block.resnets.{i}."""
      __UpperCAmelCase      :str			    =						f"""mid.block_{i+1}."""
      vae_conversion_map.append((sd_mid_res_prefix, hf_mid_res_prefix))


__UpperCAmelCase      :Tuple			    =						[
    # (stable-diffusion, HF Diffusers)
    ("norm.", "group_norm."),
    ("q.", "query."),
    ("k.", "key."),
    ("v.", "value."),
    ("proj_out.", "proj_attn."),
]




def      _a     (						_lowercase :     int  ):





     '''simple docstring'''






     return w.reshape(*w.shape      , 1      , 1  )




def      _a     (						_lowercase :     Optional[Any]  ):





     '''simple docstring'''






     __UpperCAmelCase     :			str  							=      {k: k for k in vae_state_dict.keys()}
     for k, v in mapping.items():
          for sd_part, hf_part in vae_conversion_map:
               __UpperCAmelCase     :			List[str]  							=      v.replace(_lowercase      , _lowercase  )
          __UpperCAmelCase     :			int  							=      v
     for k, v in mapping.items():
          if "attentions" in k:
               for sd_part, hf_part in vae_conversion_map_attn:
                    __UpperCAmelCase     :			List[Any]  							=      v.replace(_lowercase      , _lowercase  )
               __UpperCAmelCase     :			str  							=      v
     __UpperCAmelCase     :			List[str]  							=      {v: vae_state_dict[k] for k, v in mapping.items()}
     __UpperCAmelCase     :			List[Any]  							=      ['''q''', '''k''', '''v''', '''proj_out''']
     for k, v in new_state_dict.items():
          for weight_name in weights_to_convert:
               if F'mid.attn_1.{weight_name}.weight' in k:
                    print(F'Reshaping {k} for SD format'  )
                    __UpperCAmelCase     :			List[str]  							=      reshape_weight_for_sd(_lowercase  )
     return new_state_dict


# =========================#
# Text Encoder Conversion #
# =========================#


__UpperCAmelCase      :Union[str, Any]			    =						[
    # (stable-diffusion, HF Diffusers)
    ("resblocks.", "text_model.encoder.layers."),
    ("ln_1", "layer_norm1"),
    ("ln_2", "layer_norm2"),
    (".c_fc.", ".fc1."),
    (".c_proj.", ".fc2."),
    (".attn", ".self_attn"),
    ("ln_final.", "transformer.text_model.final_layer_norm."),
    ("token_embedding.weight", "transformer.text_model.embeddings.token_embedding.weight"),
    ("positional_embedding", "transformer.text_model.embeddings.position_embedding.weight"),
]
__UpperCAmelCase      :int			    =						{re.escape(x[1]): x[0] for x in textenc_conversion_lst}
__UpperCAmelCase      :Optional[Any]			    =						re.compile("|".join(protected.keys()))

# Ordering is from https://github.com/pytorch/pytorch/blob/master/test/cpp/api/modules.cpp
__UpperCAmelCase      :int			    =						{"q": 0, "k": 1, "v": 2}




def      _a     (						_lowercase :     Optional[Any]  ):





     '''simple docstring'''






     __UpperCAmelCase     :			Optional[Any]  							=      {}
     __UpperCAmelCase     :			int  							=      {}
     __UpperCAmelCase     :			Optional[int]  							=      {}
     for k, v in text_enc_dict.items():
          if (
              k.endswith('''.self_attn.q_proj.weight'''  )
              or k.endswith('''.self_attn.k_proj.weight'''  )
              or k.endswith('''.self_attn.v_proj.weight'''  )
          ):
               __UpperCAmelCase     :			str  							=      k[: -len('''.q_proj.weight'''  )]
               __UpperCAmelCase     :			Optional[int]  							=      k[-len('''q_proj.weight'''  )]
               if k_pre not in capture_qkv_weight:
                    __UpperCAmelCase     :			Any  							=      [None, None, None]
               __UpperCAmelCase     :			List[str]  							=      v
               continue

          if (
              k.endswith('''.self_attn.q_proj.bias'''  )
              or k.endswith('''.self_attn.k_proj.bias'''  )
              or k.endswith('''.self_attn.v_proj.bias'''  )
          ):
               __UpperCAmelCase     :			List[Any]  							=      k[: -len('''.q_proj.bias'''  )]
               __UpperCAmelCase     :			List[str]  							=      k[-len('''q_proj.bias'''  )]
               if k_pre not in capture_qkv_bias:
                    __UpperCAmelCase     :			Union[str, Any]  							=      [None, None, None]
               __UpperCAmelCase     :			List[str]  							=      v
               continue

          __UpperCAmelCase     :			Dict  							=      textenc_pattern.sub(lambda _lowercase  : protected[re.escape(m.group(0  )  )]      , _lowercase  )
          __UpperCAmelCase     :			Any  							=      v

     for k_pre, tensors in capture_qkv_weight.items():
          if None in tensors:
               raise Exception('''CORRUPTED MODEL: one of the q-k-v values for the text encoder was missing'''  )
          __UpperCAmelCase     :			int  							=      textenc_pattern.sub(lambda _lowercase  : protected[re.escape(m.group(0  )  )]      , _lowercase  )
          __UpperCAmelCase     :			Dict  							=      torch.cat(_lowercase  )

     for k_pre, tensors in capture_qkv_bias.items():
          if None in tensors:
               raise Exception('''CORRUPTED MODEL: one of the q-k-v values for the text encoder was missing'''  )
          __UpperCAmelCase     :			Optional[int]  							=      textenc_pattern.sub(lambda _lowercase  : protected[re.escape(m.group(0  )  )]      , _lowercase  )
          __UpperCAmelCase     :			str  							=      torch.cat(_lowercase  )

     return new_state_dict




def      _a     (						_lowercase :     Optional[int]  ):





     '''simple docstring'''






     return text_enc_dict


if __name__ == "__main__":
      __UpperCAmelCase      :List[str]			    =						argparse.ArgumentParser()

      parser.add_argument("--model_path", default=None, type=str, required=True, help="Path to the model to convert.")
      parser.add_argument("--checkpoint_path", default=None, type=str, required=True, help="Path to the output model.")
      parser.add_argument("--half", action="store_true", help="Save weights in half precision.")
      parser.add_argument(
          "--use_safetensors", action="store_true", help="Save weights use safetensors, default is ckpt."
      )

      __UpperCAmelCase      :List[str]			    =						parser.parse_args()

      assert args.model_path is not None, "Must provide a model path!"

      assert args.checkpoint_path is not None, "Must provide a checkpoint path!"

      # Path for safetensors
      __UpperCAmelCase      :Optional[int]			    =						osp.join(args.model_path, "unet", "diffusion_pytorch_model.safetensors")
      __UpperCAmelCase      :Any			    =						osp.join(args.model_path, "vae", "diffusion_pytorch_model.safetensors")
      __UpperCAmelCase      :int			    =						osp.join(args.model_path, "text_encoder", "model.safetensors")

      # Load models from safetensors if it exists, if it doesn't pytorch
      if osp.exists(unet_path):
            __UpperCAmelCase      :int			    =						load_file(unet_path, device="cpu")
      else:
            __UpperCAmelCase      :int			    =						osp.join(args.model_path, "unet", "diffusion_pytorch_model.bin")
            __UpperCAmelCase      :Optional[Any]			    =						torch.load(unet_path, map_location="cpu")

      if osp.exists(vae_path):
            __UpperCAmelCase      :Optional[int]			    =						load_file(vae_path, device="cpu")
      else:
            __UpperCAmelCase      :Any			    =						osp.join(args.model_path, "vae", "diffusion_pytorch_model.bin")
            __UpperCAmelCase      :Optional[Any]			    =						torch.load(vae_path, map_location="cpu")

      if osp.exists(text_enc_path):
            __UpperCAmelCase      :List[Any]			    =						load_file(text_enc_path, device="cpu")
      else:
            __UpperCAmelCase      :int			    =						osp.join(args.model_path, "text_encoder", "pytorch_model.bin")
            __UpperCAmelCase      :Optional[int]			    =						torch.load(text_enc_path, map_location="cpu")

      # Convert the UNet model
      __UpperCAmelCase      :str			    =						convert_unet_state_dict(unet_state_dict)
      __UpperCAmelCase      :int			    =						{"model.diffusion_model." + k: v for k, v in unet_state_dict.items()}

      # Convert the VAE model
      __UpperCAmelCase      :Optional[int]			    =						convert_vae_state_dict(vae_state_dict)
      __UpperCAmelCase      :List[str]			    =						{"first_stage_model." + k: v for k, v in vae_state_dict.items()}

      # Easiest way to identify v2.0 model seems to be that the text encoder (OpenCLIP) is deeper
      __UpperCAmelCase      :Optional[int]			    =						"text_model.encoder.layers.22.layer_norm2.bias" in text_enc_dict

      if is_vaa_model:
            # Need to add the tag 'transformer' in advance so we can knock it out from the final layer-norm
            __UpperCAmelCase      :List[Any]			    =						{"transformer." + k: v for k, v in text_enc_dict.items()}
            __UpperCAmelCase      :Union[str, Any]			    =						convert_text_enc_state_dict_vaa(text_enc_dict)
            __UpperCAmelCase      :List[str]			    =						{"cond_stage_model.model." + k: v for k, v in text_enc_dict.items()}
      else:
            __UpperCAmelCase      :List[Any]			    =						convert_text_enc_state_dict(text_enc_dict)
            __UpperCAmelCase      :Optional[Any]			    =						{"cond_stage_model.transformer." + k: v for k, v in text_enc_dict.items()}

      # Put together new checkpoint
      __UpperCAmelCase      :str			    =						{**unet_state_dict, **vae_state_dict, **text_enc_dict}
      if args.half:
            __UpperCAmelCase      :List[Any]			    =						{k: v.half() for k, v in state_dict.items()}

      if args.use_safetensors:
            save_file(state_dict, args.checkpoint_path)
      else:
            __UpperCAmelCase      :Dict			    =						{"state_dict": state_dict}
            torch.save(state_dict, args.checkpoint_path)



'''simple docstring'''



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

import torch
import torch.nn.functional as F
from torch import nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput
from .attention import BasicTransformerBlock
from .attention_processor import AttentionProcessor, AttnProcessor
from .embeddings import TimestepEmbedding, Timesteps
from .modeling_utils import ModelMixin



@dataclass
class      a    (		_a	):


     """simple docstring"""






     SCREAMING_SNAKE_CASE	:			torch.FloatTensor



class      a    (		_a     ,  _a	):


     """simple docstring"""






     @register_to_config
     def __init__(    self					:		str							,			snake_case					:		int = 32							,			snake_case					:		int = 64							,			snake_case					:		int = 20							,			snake_case					:		int = 768							,			snake_case					:		Tuple=77							,			snake_case					:		List[Any]=4							,			snake_case					:		float = 0.0							,			snake_case					:		str = "silu"							,			snake_case					:		Optional[str] = None							,			snake_case					:		Optional[str] = None							,			snake_case					:		Optional[str] = "linear"							,			snake_case					:		Optional[str] = "prd"							,			snake_case					:		Optional[int] = None							,			snake_case					:		Optional[int] = None							,			snake_case					:		Optional[int] = None							,			)						->							List[str]:
          super().__init__()
          __UpperCAmelCase     :			List[Any]  							=      num_attention_heads
          __UpperCAmelCase     :			List[str]  							=      attention_head_dim
          __UpperCAmelCase     :			int  							=      num_attention_heads * attention_head_dim
          __UpperCAmelCase     :			List[Any]  							=      additional_embeddings

          __UpperCAmelCase     :			Any  							=      time_embed_dim or inner_dim
          __UpperCAmelCase     :			Any  							=      embedding_proj_dim or embedding_dim
          __UpperCAmelCase     :			Union[str, Any]  							=      clip_embed_dim or embedding_dim

          __UpperCAmelCase     :			List[Any]  							=      Timesteps(snake_case							,			snake_case							,			0      )
          __UpperCAmelCase     :			Optional[Any]  							=      TimestepEmbedding(snake_case							,			snake_case							,			out_dim=snake_case							,			act_fn=snake_case      )

          __UpperCAmelCase     :			str  							=      nn.Linear(snake_case							,			snake_case      )

          if embedding_proj_norm_type is None:
               __UpperCAmelCase     :			str  							=      None
          elif embedding_proj_norm_type == "layer":
               __UpperCAmelCase     :			str  							=      nn.LayerNorm(snake_case      )
          else:
               raise ValueError(f'unsupported embedding_proj_norm_type: {embedding_proj_norm_type}'      )

          __UpperCAmelCase     :			List[Any]  							=      nn.Linear(snake_case							,			snake_case      )

          if encoder_hid_proj_type is None:
               __UpperCAmelCase     :			Union[str, Any]  							=      None
          elif encoder_hid_proj_type == "linear":
               __UpperCAmelCase     :			Any  							=      nn.Linear(snake_case							,			snake_case      )
          else:
               raise ValueError(f'unsupported encoder_hid_proj_type: {encoder_hid_proj_type}'      )

          __UpperCAmelCase     :			Dict  							=      nn.Parameter(torch.zeros(1							,			num_embeddings + additional_embeddings							,			snake_case      )      )

          if added_emb_type == "prd":
               __UpperCAmelCase     :			Any  							=      nn.Parameter(torch.zeros(1							,			1							,			snake_case      )      )
          elif added_emb_type is None:
               __UpperCAmelCase     :			List[Any]  							=      None
          else:
               raise ValueError(
                   f'`added_emb_type`: {added_emb_type} is not supported. Make sure to choose one of `\'prd\'` or `None`.'      )

          __UpperCAmelCase     :			Optional[int]  							=      nn.ModuleList(
              [
                  BasicTransformerBlock(
                      snake_case							,			snake_case							,			snake_case							,			dropout=snake_case							,			activation_fn='''gelu'''							,			attention_bias=snake_case							,			)
                  for d in range(snake_case      )
              ]      )

          if norm_in_type == "layer":
               __UpperCAmelCase     :			Tuple  							=      nn.LayerNorm(snake_case      )
          elif norm_in_type is None:
               __UpperCAmelCase     :			List[Any]  							=      None
          else:
               raise ValueError(f'Unsupported norm_in_type: {norm_in_type}.'      )

          __UpperCAmelCase     :			Dict  							=      nn.LayerNorm(snake_case      )

          __UpperCAmelCase     :			Any  							=      nn.Linear(snake_case							,			snake_case      )

          __UpperCAmelCase     :			Any  							=      torch.full(
              [num_embeddings + additional_embeddings, num_embeddings + additional_embeddings]							,			-10_000.0      )
          causal_attention_mask.triu_(1      )
          __UpperCAmelCase     :			str  							=      causal_attention_mask[None, ...]
          self.register_buffer('''causal_attention_mask'''							,			snake_case							,			persistent=snake_case      )

          __UpperCAmelCase     :			Tuple  							=      nn.Parameter(torch.zeros(1							,			snake_case      )      )
          __UpperCAmelCase     :			List[str]  							=      nn.Parameter(torch.zeros(1							,			snake_case      )      )

     @property
     # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.attn_processors
     def 			lowerCamelCase__			(    self					:		int      )						->							Dict[str, AttentionProcessor]:
          __UpperCAmelCase     :			Optional[Any]  							=      {}

          def fn_recursive_add_processors(snake_case					:		str							,			snake_case					:		torch.nn.Module							,			snake_case					:		Dict[str, AttentionProcessor]      ):
               if hasattr(snake_case							,			'''set_processor'''      ):
                    __UpperCAmelCase     :			Union[str, Any]  							=      module.processor

               for sub_name, child in module.named_children():
                    fn_recursive_add_processors(f'{name}.{sub_name}'							,			snake_case							,			snake_case      )

               return processors

          for name, module in self.named_children():
               fn_recursive_add_processors(snake_case							,			snake_case							,			snake_case      )

          return processors

     def 			lowerCamelCase__			(    self					:		Optional[Any]							,			snake_case					:		Union[AttentionProcessor, Dict[str, AttentionProcessor]]      )						->							Union[str, Any]:
          __UpperCAmelCase     :			Union[str, Any]  							=      len(self.attn_processors.keys()      )

          if isinstance(snake_case							,			snake_case      ) and len(snake_case      ) != count:
               raise ValueError(
                   f'A dict of processors was passed, but the number of processors {len(snake_case      )} does not match the'
                   f' number of attention layers: {count}. Please make sure to pass {count} processor classes.'      )

          def fn_recursive_attn_processor(snake_case					:		str							,			snake_case					:		torch.nn.Module							,			snake_case					:		int      ):
               if hasattr(snake_case							,			'''set_processor'''      ):
                    if not isinstance(snake_case							,			snake_case      ):
                         module.set_processor(snake_case      )
                    else:
                         module.set_processor(processor.pop(f'{name}.processor'      )      )

               for sub_name, child in module.named_children():
                    fn_recursive_attn_processor(f'{name}.{sub_name}'							,			snake_case							,			snake_case      )

          for name, module in self.named_children():
               fn_recursive_attn_processor(snake_case							,			snake_case							,			snake_case      )

     def 			lowerCamelCase__			(    self					:		str      )						->							Tuple:
          self.set_attn_processor(AttnProcessor()      )

     def 			lowerCamelCase__			(    self					:		Optional[Any]							,			snake_case					:		List[Any]							,			snake_case					:		Union[torch.Tensor, float, int]							,			snake_case					:		torch.FloatTensor							,			snake_case					:		Optional[torch.FloatTensor] = None							,			snake_case					:		Optional[torch.BoolTensor] = None							,			snake_case					:		bool = True							,			)						->							List[Any]:
          __UpperCAmelCase     :			Any  							=      hidden_states.shape[0]

          __UpperCAmelCase     :			Optional[int]  							=      timestep
          if not torch.is_tensor(snake_case      ):
               __UpperCAmelCase     :			str  							=      torch.tensor([timesteps]							,			dtype=torch.long							,			device=hidden_states.device      )
          elif torch.is_tensor(snake_case      ) and len(timesteps.shape      ) == 0:
               __UpperCAmelCase     :			Any  							=      timesteps[None].to(hidden_states.device      )

          # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
          __UpperCAmelCase     :			Optional[int]  							=      timesteps * torch.ones(snake_case							,			dtype=timesteps.dtype							,			device=timesteps.device      )

          __UpperCAmelCase     :			Tuple  							=      self.time_proj(snake_case      )

          # timesteps does not contain any weights and will always return f32 tensors
          # but time_embedding might be fp16, so we need to cast here.
          __UpperCAmelCase     :			int  							=      timesteps_projected.to(dtype=self.dtype      )
          __UpperCAmelCase     :			Optional[int]  							=      self.time_embedding(snake_case      )

          if self.embedding_proj_norm is not None:
               __UpperCAmelCase     :			Optional[Any]  							=      self.embedding_proj_norm(snake_case      )

          __UpperCAmelCase     :			str  							=      self.embedding_proj(snake_case      )
          if self.encoder_hidden_states_proj is not None and encoder_hidden_states is not None:
               __UpperCAmelCase     :			Dict  							=      self.encoder_hidden_states_proj(snake_case      )
          elif self.encoder_hidden_states_proj is not None and encoder_hidden_states is None:
               raise ValueError('''`encoder_hidden_states_proj` requires `encoder_hidden_states` to be set'''      )

          __UpperCAmelCase     :			Optional[int]  							=      self.proj_in(snake_case      )

          __UpperCAmelCase     :			Optional[Any]  							=      self.positional_embedding.to(hidden_states.dtype      )

          __UpperCAmelCase     :			Union[str, Any]  							=      []
          __UpperCAmelCase     :			Optional[Any]  							=      0

          if encoder_hidden_states is not None:
               additional_embeds.append(snake_case      )
               additional_embeddings_len += encoder_hidden_states.shape[1]

          if len(proj_embeddings.shape      ) == 2:
               __UpperCAmelCase     :			Optional[int]  							=      proj_embeddings[:, None, :]

          if len(hidden_states.shape      ) == 2:
               __UpperCAmelCase     :			Union[str, Any]  							=      hidden_states[:, None, :]

          __UpperCAmelCase     :			Union[str, Any]  							=      additional_embeds + [
              proj_embeddings,
              time_embeddings[:, None, :],
              hidden_states,
          ]

          if self.prd_embedding is not None:
               __UpperCAmelCase     :			Any  							=      self.prd_embedding.to(hidden_states.dtype      ).expand(snake_case							,			-1							,			-1      )
               additional_embeds.append(snake_case      )

          __UpperCAmelCase     :			Dict  							=      torch.cat(
              snake_case							,			dim=1							,			)

          # Allow positional_embedding to not include the `addtional_embeddings` and instead pad it with zeros for these additional tokens
          __UpperCAmelCase     :			str  							=      additional_embeddings_len + proj_embeddings.shape[1] + 1
          if positional_embeddings.shape[1] < hidden_states.shape[1]:
               __UpperCAmelCase     :			Union[str, Any]  							=      F.pad(
                   snake_case							,			(
                       0,
                       0,
                       additional_embeddings_len,
                       self.prd_embedding.shape[1] if self.prd_embedding is not None else 0,
                   )							,			value=0.0							,			)

          __UpperCAmelCase     :			Optional[int]  							=      hidden_states + positional_embeddings

          if attention_mask is not None:
               __UpperCAmelCase     :			List[Any]  							=      (1 - attention_mask.to(hidden_states.dtype      )) * -10_000.0
               __UpperCAmelCase     :			str  							=      F.pad(snake_case							,			(0, self.additional_embeddings)							,			value=0.0      )
               __UpperCAmelCase     :			Optional[int]  							=      (attention_mask[:, None, :] + self.causal_attention_mask).to(hidden_states.dtype      )
               __UpperCAmelCase     :			Optional[int]  							=      attention_mask.repeat_interleave(self.config.num_attention_heads							,			dim=0      )

          if self.norm_in is not None:
               __UpperCAmelCase     :			str  							=      self.norm_in(snake_case      )

          for block in self.transformer_blocks:
               __UpperCAmelCase     :			Optional[int]  							=      block(snake_case							,			attention_mask=snake_case      )

          __UpperCAmelCase     :			int  							=      self.norm_out(snake_case      )

          if self.prd_embedding is not None:
               __UpperCAmelCase     :			Optional[int]  							=      hidden_states[:, -1]
          else:
               __UpperCAmelCase     :			List[Any]  							=      hidden_states[:, additional_embeddings_len:]

          __UpperCAmelCase     :			Dict  							=      self.proj_to_clip_embeddings(snake_case      )

          if not return_dict:
               return (predicted_image_embedding,)

          return PriorTransformerOutput(predicted_image_embedding=snake_case      )






     def 			lowerCamelCase__			(    self					:		Optional[int]							,			snake_case					:		List[str]      )						->							str:
          __UpperCAmelCase     :			Dict  							=      (prior_latents * self.clip_std) + self.clip_mean
          return prior_latents