



from __future__ import annotations

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

snake_case__					:     Union[str, Any]    		=   TypeVar("""T""")
snake_case__					:     Optional[int]    		=   TypeVar("""U""")

class   _A   ( Generic[T, U]	):


						'''simple docstring'''





						def __init__(	self	:							Optional[int]							,      lowerCamelCase	:							T | None							,      lowerCamelCase	:							U | None	):






							'''simple docstring'''



							__lowercase		       =			key
							__lowercase		       =			val
							__lowercase		       =			None
							__lowercase		       =			None



						def __repr__(	self	:							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]	):


						'''simple docstring'''





						def __init__(	self	:							Dict	):






							'''simple docstring'''



							__lowercase		       =			DoubleLinkedListNode(lowerCamelCase							,      lowerCamelCase	)
							__lowercase		       =			DoubleLinkedListNode(lowerCamelCase							,      lowerCamelCase	)
							__lowercase       ,    __lowercase		       =			self.rear, self.head



						def __repr__(	self	:							Optional[Any]	):






							'''simple docstring'''



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



						def     _snake_case	(	self	:							Union[str, Any]							,      lowerCamelCase	:							DoubleLinkedListNode[T, U]	):






							'''simple docstring'''



							__lowercase		       =			self.rear.prev

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

							__lowercase		       =			node
							__lowercase		       =			previous
							__lowercase		       =			node
							__lowercase		       =			self.rear



						def     _snake_case	(	self	:							Optional[int]							,      lowerCamelCase	:							DoubleLinkedListNode[T, U]	):






							'''simple docstring'''



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

							__lowercase		       =			node.next
							__lowercase		       =			node.prev
							__lowercase		       =			None
							__lowercase		       =			None
							return node

class   _A   ( Generic[T, U]	):


						'''simple docstring'''





						_snake_case			:     dict[Callable[[T], U], LRUCache[T, U]]      							=      {}



						def __init__(	self	:							List[Any]							,      lowerCamelCase	:							int	):






							'''simple docstring'''



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



						def __repr__(	self	:							Optional[Any]	):






							'''simple docstring'''



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



						def __contains__(	self	:							Dict							,      lowerCamelCase	:							T	):






							'''simple docstring'''



							return key in self.cache



						def     _snake_case	(	self	:							List[Any]							,      lowerCamelCase	:							T	):






							'''simple docstring'''



							if key in self.cache:
								self.hits += 1
								__lowercase		       =			self.cache[key]
								__lowercase		       =			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(lowerCamelCase	)
								return node.val
							self.miss += 1
							return None



						def     _snake_case	(	self	:							Union[str, Any]							,      lowerCamelCase	:							T							,      lowerCamelCase	:							U	):






							'''simple docstring'''



							if key not in self.cache:
								if self.num_keys >= self.capacity:
									# delete first node (oldest) when over capacity
									__lowercase		       =			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(lowerCamelCase	) 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
								__lowercase		       =			DoubleLinkedListNode(lowerCamelCase							,      lowerCamelCase	)
								self.list.add(self.cache[key]	)
								self.num_keys += 1

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



						@classmethod
						def     _snake_case	(	cls	:							Union[str, Any]							,      lowerCamelCase	:							int = 128	):






							'''simple docstring'''



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

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

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

								setattr(lowerCamelCase							,      "cache_info"							,      lowerCamelCase	)  # noqa: B010

								return cache_decorator_wrapper

							return cache_decorator_inner


if __name__ == "__main__":
	import doctest

	doctest.testmod()








import torch
from torch import nn

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

class   _A   ( _lowercase  ,						_lowercase	):


						'''simple docstring'''





						@register_to_config
						def __init__(	self	:							Optional[Any]							,      *,
						    lowerCamelCase	:							int = 4							,      lowerCamelCase	:							int = 768							,      lowerCamelCase	:							int							,      lowerCamelCase	:							Optional[int]							,      ):






							'''simple docstring'''



							super().__init__()

							__lowercase		       =			nn.Parameter(torch.zeros(lowerCamelCase	)	)

							# parameters for additional clip time embeddings
							__lowercase		       =			nn.Linear(lowerCamelCase							,      lowerCamelCase	)
							__lowercase		       =			nn.Linear(lowerCamelCase							,      lowerCamelCase	)

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



						def     _snake_case	(	self	:							Union[str, Any]							,      *, lowerCamelCase	:							Any							,      lowerCamelCase	:							Tuple							,      lowerCamelCase	:							Optional[int]							,      lowerCamelCase	:							Tuple	):






							'''simple docstring'''



							if do_classifier_free_guidance:
								# Add the classifier free guidance embeddings to the image embeddings
								__lowercase		       =			image_embeddings.shape[0]
								__lowercase		       =			self.learned_classifier_free_guidance_embeddings.unsqueeze(0	)
								__lowercase		       =			classifier_free_guidance_embeddings.expand(
								    lowerCamelCase							,      -1	)
								__lowercase		       =			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]

							__lowercase		       =			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, ...
							__lowercase		       =			self.embedding_proj(lowerCamelCase	)
							__lowercase		       =			self.clip_image_embeddings_project_to_time_embeddings(lowerCamelCase	)
							__lowercase		       =			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"
							__lowercase		       =			self.clip_extra_context_tokens_proj(lowerCamelCase	)
							__lowercase		       =			clip_extra_context_tokens.reshape(lowerCamelCase							,      -1							,      self.clip_extra_context_tokens	)
							__lowercase		       =			clip_extra_context_tokens.permute(0							,      2							,      1	)

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

							return text_encoder_hidden_states, additive_clip_time_embeddings



