'''simple docstring'''




from __future__ import annotations

from collections.abc import Callable
from typing import Generic, TypeVar

__UpperCAmelCase												=TypeVar("T")
__UpperCAmelCase												=TypeVar("U")




class    a__				(       Generic[T, U]			):






     def __init__(			self		:		Dict  ,    a		:		T | None  ,    a		:		U | None   ):





         """simple docstring"""



         __lowerCamelCase          =  key
         __lowerCamelCase          =  val
         __lowerCamelCase          =  None
         __lowerCamelCase          =  None






     def __repr__(			self		:		Optional[Any]   ):





         """simple docstring"""



         return (
             f"""Node: key: {self.key}, val: {self.val}, """
             f"""has next: {bool(self.next   )}, has prev: {bool(self.prev   )}"""
         )




class    a__				(       Generic[T, U]			):






     def __init__(			self		:		Optional[Any]   ):





         """simple docstring"""



         __lowerCamelCase          =  DoubleLinkedListNode(a  ,    a   )
         __lowerCamelCase          =  DoubleLinkedListNode(a  ,    a   )
         __lowerCamelCase     ,  __lowerCamelCase          =  self.rear, self.head






     def __repr__(			self		:		Dict   ):





         """simple docstring"""



         __lowerCamelCase          =  ['''DoubleLinkedList''']
         __lowerCamelCase          =  self.head
         while node.next is not None:
             rep.append(str(a   )   )
             __lowerCamelCase          =  node.next
         rep.append(str(self.rear   )   )
         return ",\n    ".join(a   )






     def        SCREAMING_SNAKE_CASE__       (			self		:		Any  ,    a		:		DoubleLinkedListNode[T, U]   ):





         """simple docstring"""



         __lowerCamelCase          =  self.rear.prev

         # All nodes other than self.head are guaranteed to have non-None previous
         assert previous is not None

         __lowerCamelCase          =  node
         __lowerCamelCase          =  previous
         __lowerCamelCase          =  node
         __lowerCamelCase          =  self.rear






     def        SCREAMING_SNAKE_CASE__       (			self		:		Tuple  ,    a		:		DoubleLinkedListNode[T, U]   ):





         """simple docstring"""



         if node.prev is None or node.next is None:
             return None

         __lowerCamelCase          =  node.next
         __lowerCamelCase          =  node.prev
         __lowerCamelCase          =  None
         __lowerCamelCase          =  None
         return node




class    a__				(       Generic[T, U]			):
     lowerCamelCase  :					dict[Callable[[T], U], LRUCache[T, U]]							={}






     def __init__(			self		:		Tuple  ,    a		:		int   ):





         """simple docstring"""



         __lowerCamelCase          =  DoubleLinkedList()
         __lowerCamelCase          =  capacity
         __lowerCamelCase          =  0
         __lowerCamelCase          =  0
         __lowerCamelCase          =  0
         __lowerCamelCase          =  {}






     def __repr__(			self		:		Optional[int]   ):





         """simple docstring"""



         return (
             f"""CacheInfo(hits={self.hits}, misses={self.miss}, """
             f"""capacity={self.capacity}, current size={self.num_keys})"""
         )






     def __contains__(			self		:		Optional[Any]  ,    a		:		T   ):





         """simple docstring"""



         return key in self.cache






     def        SCREAMING_SNAKE_CASE__       (			self		:		Tuple  ,    a		:		T   ):





         """simple docstring"""



         if key in self.cache:
             self.hits += 1
             __lowerCamelCase          =  self.cache[key]
             __lowerCamelCase          =  self.list.remove(self.cache[key]   )
             assert node == value_node

             # node is guaranteed not None because it is in self.cache
             assert node is not None
             self.list.add(a   )
             return node.val
         self.miss += 1
         return None






     def        SCREAMING_SNAKE_CASE__       (			self		:		int  ,    a		:		T  ,    a		:		U   ):





         """simple docstring"""



         if key not in self.cache:
             if self.num_keys >= self.capacity:
                 # delete first node (oldest) when over capacity
                 __lowerCamelCase          =  self.list.head.next

                 # guaranteed to have a non-None first node when num_keys > 0
                 # explain to type checker via assertions
                 assert first_node is not None
                 assert first_node.key is not None
                 assert (
                     self.list.remove(a   ) is not None
                 )  # node guaranteed to be in list assert node.key is not None

                 del self.cache[first_node.key]
                 self.num_keys -= 1
             __lowerCamelCase          =  DoubleLinkedListNode(a  ,    a   )
             self.list.add(self.cache[key]   )
             self.num_keys += 1

         else:
             # bump node to the end of the list, update value
             __lowerCamelCase          =  self.list.remove(self.cache[key]   )
             assert node is not None  # node guaranteed to be in list
             __lowerCamelCase          =  value
             self.list.add(a   )






     @classmethod
     def        SCREAMING_SNAKE_CASE__       (			cls		:		Any  ,    a		:		int = 1_28   ):





         """simple docstring"""



         def cache_decorator_inner(a		:		Callable[[T], U]   ) -> Callable[..., U]:
             def cache_decorator_wrapper(*a		:		T   ) -> U:
                 if func not in cls.decorator_function_to_instance_map:
                     __lowerCamelCase          =  LRUCache(a   )

                 __lowerCamelCase          =  cls.decorator_function_to_instance_map[func].get(args[0]   )
                 if result is None:
                     __lowerCamelCase          =  func(*a   )
                     cls.decorator_function_to_instance_map[func].put(args[0]  ,    a   )
                 return result

             def cache_info() -> LRUCache[T, U]:
                 return cls.decorator_function_to_instance_map[func]

             setattr(a  ,    '''cache_info'''  ,    a   )  # noqa: B010

             return cache_decorator_wrapper

         return cache_decorator_inner


if __name__ == "__main__":
  import doctest

  doctest.testmod()

'''simple docstring'''




import torch
from torch import nn

from ...configuration_utils import ConfigMixin, register_to_config
from ...models import ModelMixin




class    a__				(       UpperCAmelCase__							,       UpperCAmelCase__			):






     @register_to_config
     def __init__(			self		:		Any  ,    *,
         a		:		int = 4  ,    a		:		int = 7_68  ,    a		:		int  ,    a		:		Optional[int]  ,    ):





         """simple docstring"""



         super().__init__()

         __lowerCamelCase          =  nn.Parameter(torch.zeros(a   )   )

         # parameters for additional clip time embeddings
         __lowerCamelCase          =  nn.Linear(a  ,    a   )
         __lowerCamelCase          =  nn.Linear(a  ,    a   )

         # parameters for encoder hidden states
         __lowerCamelCase          =  clip_extra_context_tokens
         __lowerCamelCase          =  nn.Linear(
             a  ,    self.clip_extra_context_tokens * cross_attention_dim   )
         __lowerCamelCase          =  nn.Linear(a  ,    a   )
         __lowerCamelCase          =  nn.LayerNorm(a   )






     def        SCREAMING_SNAKE_CASE__       (			self		:		int  ,    *, a		:		Union[str, Any]  ,    a		:		Any  ,    a		:		str  ,    a		:		Any   ):





         """simple docstring"""



         if do_classifier_free_guidance:
             # Add the classifier free guidance embeddings to the image embeddings
             __lowerCamelCase          =  image_embeddings.shape[0]
             __lowerCamelCase          =  self.learned_classifier_free_guidance_embeddings.unsqueeze(0   )
             __lowerCamelCase          =  classifier_free_guidance_embeddings.expand(
                 a  ,    -1   )
             __lowerCamelCase          =  torch.cat([classifier_free_guidance_embeddings, image_embeddings]  ,    dim=0   )

         # The image embeddings batch size and the text embeddings batch size are equal
         assert image_embeddings.shape[0] == prompt_embeds.shape[0]

         __lowerCamelCase          =  prompt_embeds.shape[0]

         # "Specifically, we modify the architecture described in Nichol et al. (2021) by projecting and
         # adding CLIP embeddings to the existing timestep embedding, ...
         __lowerCamelCase          =  self.embedding_proj(a   )
         __lowerCamelCase          =  self.clip_image_embeddings_project_to_time_embeddings(a   )
         __lowerCamelCase          =  time_projected_image_embeddings + time_projected_prompt_embeds

         # ... and by projecting CLIP embeddings into four
         # extra tokens of context that are concatenated to the sequence of outputs from the GLIDE text encoder"
         __lowerCamelCase          =  self.clip_extra_context_tokens_proj(a   )
         __lowerCamelCase          =  clip_extra_context_tokens.reshape(a  ,    -1  ,    self.clip_extra_context_tokens   )
         __lowerCamelCase          =  clip_extra_context_tokens.permute(0  ,    2  ,    1   )

         __lowerCamelCase          =  self.encoder_hidden_states_proj(a   )
         __lowerCamelCase          =  self.text_encoder_hidden_states_norm(a   )
         __lowerCamelCase          =  torch.cat([clip_extra_context_tokens, text_encoder_hidden_states]  ,    dim=1   )

         return text_encoder_hidden_states, additive_clip_time_embeddings
