import random

import torch
from huggingface_hub import HfApi

from diffusers import UNetaDModel


lowerCAmelCase_						   =    HfApi()

lowerCAmelCase_						   =    {}
# fmt: off
lowerCAmelCase_						   =    torch.tensor([
    -0.7515, -1.6883, 0.2420, 0.0300, 0.6347, 1.3433, -1.1743, -3.7467,
    1.2342, -2.2485, 0.4636, 0.8076, -0.7991, 0.3969, 0.8498, 0.9189,
    -1.8887, -3.3522, 0.7639, 0.2040, 0.6271, -2.7148, -1.6316, 3.0839,
    0.3186, 0.2721, -0.9759, -1.2461, 2.6257, 1.3557
])
lowerCAmelCase_						   =    torch.tensor([
    -2.3639, -2.5344, 0.0054, -0.6674, 1.5990, 1.0158, 0.3124, -2.1436,
    1.8795, -2.5429, -0.1566, -0.3973, 1.2490, 2.6447, 1.2283, -0.5208,
    -2.8154, -3.5119, 2.3838, 1.2033, 1.7201, -2.1256, -1.4576, 2.7948,
    2.4204, -0.9752, -1.2546, 0.8027, 3.2758, 3.1365
])
lowerCAmelCase_						   =    torch.tensor([
    -0.6531, -0.6891, -0.3172, -0.5375, -0.9140, -0.5367, -0.1175, -0.7869,
    -0.3808, -0.4513, -0.2098, -0.0083, 0.3183, 0.5140, 0.2247, -0.1304,
    -0.1302, -0.2802, -0.2084, -0.2025, -0.4967, -0.4873, -0.0861, 0.6925,
    0.0250, 0.1290, -0.1543, 0.6316, 1.0460, 1.4943
])
lowerCAmelCase_						   =    torch.tensor([
    0.0911, 0.1107, 0.0182, 0.0435, -0.0805, -0.0608, 0.0381, 0.2172,
    -0.0280, 0.1327, -0.0299, -0.0255, -0.0050, -0.1170, -0.1046, 0.0309,
    0.1367, 0.1728, -0.0533, -0.0748, -0.0534, 0.1624, 0.0384, -0.1805,
    -0.0707, 0.0642, 0.0220, -0.0134, -0.1333, -0.1505
])
lowerCAmelCase_						   =    torch.tensor([
    0.1321, 0.1337, 0.0440, 0.0622, -0.0591, -0.0370, 0.0503, 0.2133,
    -0.0177, 0.1415, -0.0116, -0.0112, 0.0044, -0.0980, -0.0789, 0.0395,
    0.1502, 0.1785, -0.0488, -0.0514, -0.0404, 0.1539, 0.0454, -0.1559,
    -0.0665, 0.0659, 0.0383, -0.0005, -0.1266, -0.1386
])
lowerCAmelCase_						   =    torch.tensor([
    0.1154, 0.1218, 0.0307, 0.0526, -0.0711, -0.0541, 0.0366, 0.2078,
    -0.0267, 0.1317, -0.0226, -0.0193, -0.0014, -0.1055, -0.0902, 0.0330,
    0.1391, 0.1709, -0.0562, -0.0693, -0.0560, 0.1482, 0.0381, -0.1683,
    -0.0681, 0.0661, 0.0331, -0.0046, -0.1268, -0.1431
])
lowerCAmelCase_						   =    torch.tensor([
    0.1192, 0.1240, 0.0414, 0.0606, -0.0557, -0.0412, 0.0430, 0.2042,
    -0.0200, 0.1385, -0.0115, -0.0132, 0.0017, -0.0965, -0.0802, 0.0398,
    0.1433, 0.1747, -0.0458, -0.0533, -0.0407, 0.1545, 0.0419, -0.1574,
    -0.0645, 0.0626, 0.0341, -0.0010, -0.1199, -0.1390
])
lowerCAmelCase_						   =    torch.tensor([
    0.1075, 0.1074, 0.0205, 0.0431, -0.0774, -0.0607, 0.0298, 0.2042,
    -0.0320, 0.1267, -0.0281, -0.0250, -0.0064, -0.1091, -0.0946, 0.0290,
    0.1328, 0.1650, -0.0580, -0.0738, -0.0586, 0.1440, 0.0337, -0.1746,
    -0.0712, 0.0605, 0.0250, -0.0099, -0.1316, -0.1473
])
lowerCAmelCase_						   =    torch.tensor([
    -1.4572, -2.0481, -0.0414, -0.6005, 1.4136, 0.5848, 0.4028, -2.7330,
    1.2212, -2.1228, 0.2155, 0.4039, 0.7662, 2.0535, 0.7477, -0.3243,
    -2.1758, -2.7648, 1.6947, 0.7026, 1.2338, -1.6078, -0.8682, 2.2810,
    1.8574, -0.5718, -0.5586, -0.0186, 2.3415, 2.1251])
lowerCAmelCase_						   =    torch.tensor([
    -1.3690, -1.9720, -0.4090, -0.6966, 1.4660, 0.9938, -0.1385, -2.7324,
    0.7736, -1.8917, 0.2923, 0.4293, 0.1693, 1.4112, 1.1887, -0.3181,
    -2.2160, -2.6381, 1.3170, 0.8163, 0.9240, -1.6544, -0.6099, 2.5259,
    1.6430, -0.9090, -0.9392, -0.0126, 2.4268, 2.3266
])
lowerCAmelCase_						   =    torch.tensor([
    -1.3525, -1.9628, -0.3956, -0.6860, 1.4664, 1.0014, -0.1259, -2.7212,
    0.7772, -1.8811, 0.2996, 0.4388, 0.1704, 1.4029, 1.1701, -0.3027,
    -2.2053, -2.6287, 1.3350, 0.8131, 0.9274, -1.6292, -0.6098, 2.5131,
    1.6505, -0.8958, -0.9298, -0.0151, 2.4257, 2.3355
])
lowerCAmelCase_						   =    torch.tensor([
    -2.0585, -2.7897, -0.2850, -0.8940, 1.9052, 0.5702, 0.6345, -3.8959,
    1.5932, -3.2319, 0.1974, 0.0287, 1.7566, 2.6543, 0.8387, -0.5351,
    -3.2736, -4.3375, 2.9029, 1.6390, 1.4640, -2.1701, -1.9013, 2.9341,
    3.4981, -0.6255, -1.1644, -0.1591, 3.7097, 3.2066
])
lowerCAmelCase_						   =    torch.tensor([
    -2.3139, -2.5594, -0.0197, -0.6785, 1.7001, 1.1606, 0.3075, -2.1740,
    1.8071, -2.5630, -0.0926, -0.3811, 1.2116, 2.6246, 1.2731, -0.5398,
    -2.8153, -3.6140, 2.3893, 1.3262, 1.6258, -2.1856, -1.3267, 2.8395,
    2.3779, -1.0623, -1.2468, 0.8959, 3.3367, 3.2243
])
lowerCAmelCase_						   =    torch.tensor([
    -2.0628, -2.7667, -0.2089, -0.8263, 2.0539, 0.5992, 0.6495, -3.8336,
    1.6025, -3.2817, 0.1721, -0.0633, 1.7516, 2.7039, 0.8100, -0.5908,
    -3.2113, -4.4343, 2.9257, 1.3632, 1.5562, -2.1489, -1.9894, 3.0560,
    3.3396, -0.7328, -1.0417, 0.0383, 3.7093, 3.2343
])
lowerCAmelCase_						   =    torch.tensor([
    -1.4574, -2.0569, -0.0473, -0.6117, 1.4018, 0.5769, 0.4129, -2.7344,
    1.2241, -2.1397, 0.2000, 0.3937, 0.7616, 2.0453, 0.7324, -0.3391,
    -2.1746, -2.7744, 1.6963, 0.6921, 1.2187, -1.6172, -0.8877, 2.2439,
    1.8471, -0.5839, -0.5605, -0.0464, 2.3250, 2.1219
])
# fmt: on

lowerCAmelCase_						   =    api.list_models(filter="""diffusers""")
for mod in models:
      if "google" in mod.author or mod.modelId == "CompVis/ldm-celebahq-256":
            lowerCAmelCase_						   =    '''/home/patrick/google_checkpoints/''' + mod.modelId.split("""/""")[-1]

            print(F"""Started running {mod.modelId}!!!""")

            if mod.modelId.startswith("""CompVis"""):
                  lowerCAmelCase_						   =    UNetaDModel.from_pretrained(local_checkpoint, subfolder="""unet""")
            else:
                  lowerCAmelCase_						   =    UNetaDModel.from_pretrained(local_checkpoint)

            torch.manual_seed(0)
            random.seed(0)

            lowerCAmelCase_						   =    torch.randn(1, model.config.in_channels, model.config.sample_size, model.config.sample_size)
            lowerCAmelCase_						   =    torch.tensor([10] * noise.shape[0])
            with torch.no_grad():
                  lowerCAmelCase_						   =    model(noise, time_step).sample

            assert torch.allclose(
                logits[0, 0, 0, :30], results["""_""".join("""_""".join(mod.modelId.split("""/""")).split("""-"""))], atol=1E-3
            )
            print(F"""{mod.modelId} has passed successfully!!!""")











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        _lowerCAmelCase    (			UpperCAmelCase_							):


  '''simple docstring'''
  a_							:       torch.FloatTensor




class        _lowerCAmelCase    (			UpperCAmelCase_							,    UpperCAmelCase_							):


  '''simple docstring'''







  @register_to_config
  def __init__(      self      :					str      ,				UpperCamelCase      :					int = 32      ,				UpperCamelCase      :					int = 64      ,				UpperCamelCase      :					int = 20      ,				UpperCamelCase      :					int = 7_68      ,				UpperCamelCase      :					Optional[int]=77      ,				UpperCamelCase      :					int=4      ,				UpperCamelCase      :					float = 0.0      ,				UpperCamelCase      :					str = "silu"      ,				UpperCamelCase      :					Optional[str] = None      ,				UpperCamelCase      :					Optional[str] = None      ,				UpperCamelCase      :					Optional[str] = "linear"      ,				UpperCamelCase      :					Optional[str] = "prd"      ,				UpperCamelCase      :					Optional[int] = None      ,				UpperCamelCase      :					Optional[int] = None      ,				UpperCamelCase      :					Optional[int] = None      ,				):


    '''simple docstring'''





    super().__init__()
    _snake_case   :				str														=	num_attention_heads
    _snake_case   :				Optional[int]														=	attention_head_dim
    _snake_case   :				Any														=	num_attention_heads * attention_head_dim
    _snake_case   :				List[Any]														=	additional_embeddings

    _snake_case   :				List[str]														=	time_embed_dim or inner_dim
    _snake_case   :				int														=	embedding_proj_dim or embedding_dim
    _snake_case   :				List[Any]														=	clip_embed_dim or embedding_dim

    _snake_case   :				Optional[Any]														=	Timesteps(UpperCamelCase      ,				UpperCamelCase      ,				0							)
    _snake_case   :				List[Any]														=	TimestepEmbedding(UpperCamelCase      ,				UpperCamelCase      ,				out_dim=UpperCamelCase      ,				act_fn=UpperCamelCase							)

    _snake_case   :				Optional[int]														=	nn.Linear(UpperCamelCase      ,				UpperCamelCase							)

    if embedding_proj_norm_type is None:
      _snake_case   :				str														=	None
    elif embedding_proj_norm_type == "layer":
      _snake_case   :				List[Any]														=	nn.LayerNorm(UpperCamelCase							)
    else:
      raise ValueError(f"""unsupported embedding_proj_norm_type: {embedding_proj_norm_type}"""							)

    _snake_case   :				str														=	nn.Linear(UpperCamelCase      ,				UpperCamelCase							)

    if encoder_hid_proj_type is None:
      _snake_case   :				Any														=	None
    elif encoder_hid_proj_type == "linear":
      _snake_case   :				Optional[int]														=	nn.Linear(UpperCamelCase      ,				UpperCamelCase							)
    else:
      raise ValueError(f"""unsupported encoder_hid_proj_type: {encoder_hid_proj_type}"""							)

    _snake_case   :				List[str]														=	nn.Parameter(torch.zeros(1      ,				num_embeddings + additional_embeddings      ,				UpperCamelCase							)							)

    if added_emb_type == "prd":
      _snake_case   :				str														=	nn.Parameter(torch.zeros(1      ,				1      ,				UpperCamelCase							)							)
    elif added_emb_type is None:
      _snake_case   :				Dict														=	None
    else:
      raise ValueError(
          f"""`added_emb_type`: {added_emb_type} is not supported. Make sure to choose one of `'prd'` or `None`."""							)

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

    if norm_in_type == "layer":
      _snake_case   :				Optional[int]														=	nn.LayerNorm(UpperCamelCase							)
    elif norm_in_type is None:
      _snake_case   :				Optional[Any]														=	None
    else:
      raise ValueError(f"""Unsupported norm_in_type: {norm_in_type}."""							)

    _snake_case   :				Optional[Any]														=	nn.LayerNorm(UpperCamelCase							)

    _snake_case   :				Union[str, Any]														=	nn.Linear(UpperCamelCase      ,				UpperCamelCase							)

    _snake_case   :				List[Any]														=	torch.full(
        [num_embeddings + additional_embeddings, num_embeddings + additional_embeddings]      ,				-1_00_00.0							)
    causal_attention_mask.triu_(1							)
    _snake_case   :				Optional[Any]														=	causal_attention_mask[None, ...]
    self.register_buffer('causal_attention_mask'      ,				UpperCamelCase      ,				persistent=UpperCamelCase							)

    _snake_case   :				str														=	nn.Parameter(torch.zeros(1      ,				UpperCamelCase							)							)
    _snake_case   :				List[str]														=	nn.Parameter(torch.zeros(1      ,				UpperCamelCase							)							)







  @property
  # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.attn_processors
  def UpperCamelCase_ (      self      :					Optional[Any]							):


    '''simple docstring'''





    _snake_case   :				Optional[Any]														=	{}

    def fn_recursive_add_processors(UpperCamelCase      :					str      ,				UpperCamelCase      :					torch.nn.Module      ,				UpperCamelCase      :					Dict[str, AttentionProcessor]							):
      if hasattr(UpperCamelCase      ,				'set_processor'							):
        _snake_case   :				Tuple														=	module.processor

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

      return processors

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

    return processors







  def UpperCamelCase_ (      self      :					List[Any]      ,				UpperCamelCase      :					Union[AttentionProcessor, Dict[str, AttentionProcessor]]							):


    '''simple docstring'''





    _snake_case   :				Optional[int]														=	len(self.attn_processors.keys()							)

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

    def fn_recursive_attn_processor(UpperCamelCase      :					str      ,				UpperCamelCase      :					torch.nn.Module      ,				UpperCamelCase      :					Union[str, Any]							):
      if hasattr(UpperCamelCase      ,				'set_processor'							):
        if not isinstance(UpperCamelCase      ,				UpperCamelCase							):
          module.set_processor(UpperCamelCase							)
        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}"""      ,				UpperCamelCase      ,				UpperCamelCase							)

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







  def UpperCamelCase_ (      self      :					Optional[Any]							):


    '''simple docstring'''





    self.set_attn_processor(AttnProcessor()							)







  def UpperCamelCase_ (      self      :					Optional[Any]      ,				UpperCamelCase      :					Any      ,				UpperCamelCase      :					Union[torch.Tensor, float, int]      ,				UpperCamelCase      :					torch.FloatTensor      ,				UpperCamelCase      :					Optional[torch.FloatTensor] = None      ,				UpperCamelCase      :					Optional[torch.BoolTensor] = None      ,				UpperCamelCase      :					bool = True      ,				):


    '''simple docstring'''





    _snake_case   :				Dict														=	hidden_states.shape[0]

    _snake_case   :				str														=	timestep
    if not torch.is_tensor(UpperCamelCase							):
      _snake_case   :				Dict														=	torch.tensor([timesteps]      ,				dtype=torch.long      ,				device=hidden_states.device							)
    elif torch.is_tensor(UpperCamelCase							) and len(timesteps.shape							) == 0:
      _snake_case   :				Tuple														=	timesteps[None].to(hidden_states.device							)

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

    _snake_case   :				Union[str, Any]														=	self.time_proj(UpperCamelCase							)

    # timesteps does not contain any weights and will always return f32 tensors
    # but time_embedding might be fp16, so we need to cast here.
    _snake_case   :				Tuple														=	timesteps_projected.to(dtype=self.dtype							)
    _snake_case   :				List[Any]														=	self.time_embedding(UpperCamelCase							)

    if self.embedding_proj_norm is not None:
      _snake_case   :				Optional[Any]														=	self.embedding_proj_norm(UpperCamelCase							)

    _snake_case   :				Union[str, Any]														=	self.embedding_proj(UpperCamelCase							)
    if self.encoder_hidden_states_proj is not None and encoder_hidden_states is not None:
      _snake_case   :				Dict														=	self.encoder_hidden_states_proj(UpperCamelCase							)
    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'							)

    _snake_case   :				str														=	self.proj_in(UpperCamelCase							)

    _snake_case   :				int														=	self.positional_embedding.to(hidden_states.dtype							)

    _snake_case   :				Optional[int]														=	[]
    _snake_case   :				List[Any]														=	0

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

    if len(proj_embeddings.shape							) == 2:
      _snake_case   :				str														=	proj_embeddings[:, None, :]

    if len(hidden_states.shape							) == 2:
      _snake_case   :				str														=	hidden_states[:, None, :]

    _snake_case   :				str														=	additional_embeds + [
        proj_embeddings,
        time_embeddings[:, None, :],
        hidden_states,
    ]

    if self.prd_embedding is not None:
      _snake_case   :				int														=	self.prd_embedding.to(hidden_states.dtype							).expand(UpperCamelCase      ,				-1      ,				-1							)
      additional_embeds.append(UpperCamelCase							)

    _snake_case   :				Optional[int]														=	torch.cat(
        UpperCamelCase      ,				dim=1      ,				)

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

    _snake_case   :				Optional[Any]														=	hidden_states + positional_embeddings

    if attention_mask is not None:
      _snake_case   :				Any														=	(1 - attention_mask.to(hidden_states.dtype							)) * -1_00_00.0
      _snake_case   :				Tuple														=	F.pad(UpperCamelCase      ,				(0, self.additional_embeddings)      ,				value=0.0							)
      _snake_case   :				int														=	(attention_mask[:, None, :] + self.causal_attention_mask).to(hidden_states.dtype							)
      _snake_case   :				str														=	attention_mask.repeat_interleave(self.config.num_attention_heads      ,				dim=0							)

    if self.norm_in is not None:
      _snake_case   :				Tuple														=	self.norm_in(UpperCamelCase							)

    for block in self.transformer_blocks:
      _snake_case   :				Any														=	block(UpperCamelCase      ,				attention_mask=UpperCamelCase							)

    _snake_case   :				Dict														=	self.norm_out(UpperCamelCase							)

    if self.prd_embedding is not None:
      _snake_case   :				str														=	hidden_states[:, -1]
    else:
      _snake_case   :				Any														=	hidden_states[:, additional_embeddings_len:]

    _snake_case   :				List[Any]														=	self.proj_to_clip_embeddings(UpperCamelCase							)

    if not return_dict:
      return (predicted_image_embedding,)

    return PriorTransformerOutput(predicted_image_embedding=UpperCamelCase							)







  def UpperCamelCase_ (      self      :					Tuple      ,				UpperCamelCase      :					Union[str, Any]							):


    '''simple docstring'''





    _snake_case   :				List[Any]														=	(prior_latents * self.clip_std) + self.clip_mean
    return prior_latents




