





'''simple docstring'''






from typing import Optional, Tuple, Union

import flax
import flax.linen as nn
import jax
import jax.numpy as jnp
from flax.core.frozen_dict import FrozenDict

from ..configuration_utils import ConfigMixin, flax_register_to_config
from ..utils import BaseOutput
from .embeddings_flax import FlaxTimestepEmbedding, FlaxTimesteps
from .modeling_flax_utils import FlaxModelMixin
from .unet_ad_blocks_flax import (
    FlaxCrossAttnDownBlockaD,
    FlaxDownBlockaD,
    FlaxUNetMidBlockaDCrossAttn,
)







@flax.struct.dataclass
class 			lowerCamelCase					(						__UpperCAmelCase  ):
					_SCREAMING_SNAKE_CASE							   =    42
					_SCREAMING_SNAKE_CASE							   =    42







class 			lowerCamelCase					(						nn.Module  ):
					_SCREAMING_SNAKE_CASE							   =    42
					_SCREAMING_SNAKE_CASE							   =    (16, 32, 96, 256)
					_SCREAMING_SNAKE_CASE							   =    jnp.floataa


					def    SCREAMING_SNAKE_CASE_		(					self	:  List[Any]       ):





											'''simple docstring'''





											_snake_case:   int  					=		nn.Conv(
											    self.block_out_channels[0]       ,     kernel_size=(3, 3)       ,     padding=((1, 1), (1, 1))       ,     dtype=self.dtype       ,     )

											_snake_case:   List[str]  					=		[]
											for i in range(len(self.block_out_channels       ) - 1       ):
																	_snake_case:   Optional[int]  					=		self.block_out_channels[i]
																	_snake_case:   int  					=		self.block_out_channels[i + 1]
																	_snake_case:   Any  					=		nn.Conv(
																	    __snake_case       ,     kernel_size=(3, 3)       ,     padding=((1, 1), (1, 1))       ,     dtype=self.dtype       ,     )
																	blocks.append(__snake_case       )
																	_snake_case:   Dict  					=		nn.Conv(
																	    __snake_case       ,     kernel_size=(3, 3)       ,     strides=(2, 2)       ,     padding=((1, 1), (1, 1))       ,     dtype=self.dtype       ,     )
																	blocks.append(__snake_case       )
											_snake_case:   List[Any]  					=		blocks

											_snake_case:   Tuple  					=		nn.Conv(
											    self.conditioning_embedding_channels       ,     kernel_size=(3, 3)       ,     padding=((1, 1), (1, 1))       ,     kernel_init=nn.initializers.zeros_init()       ,     bias_init=nn.initializers.zeros_init()       ,     dtype=self.dtype       ,     )




					def __call__(					self	:  Union[str, Any]       ,     __snake_case	:  Tuple       ):





											'''simple docstring'''





											_snake_case:   int  					=		self.conv_in(__snake_case       )
											_snake_case:   Any  					=		nn.silu(__snake_case       )

											for block in self.blocks:
																	_snake_case:   Any  					=		block(__snake_case       )
																	_snake_case:   Any  					=		nn.silu(__snake_case       )

											_snake_case:   Tuple  					=		self.conv_out(__snake_case       )

											return embedding







@flax_register_to_config
class 			lowerCamelCase					(						nn.Module     ,							__UpperCAmelCase     ,							__UpperCAmelCase  ):
					_SCREAMING_SNAKE_CASE							   =    32
					_SCREAMING_SNAKE_CASE							   =    4
					_SCREAMING_SNAKE_CASE							   =    (
					    "CrossAttnDownBlock2D",
					    "CrossAttnDownBlock2D",
					    "CrossAttnDownBlock2D",
					    "DownBlock2D",
					)
					_SCREAMING_SNAKE_CASE							   =    False
					_SCREAMING_SNAKE_CASE							   =    (320, 640, 1280, 1280)
					_SCREAMING_SNAKE_CASE							   =    2
					_SCREAMING_SNAKE_CASE							   =    8
					_SCREAMING_SNAKE_CASE							   =    None
					_SCREAMING_SNAKE_CASE							   =    1280
					_SCREAMING_SNAKE_CASE							   =    0.0
					_SCREAMING_SNAKE_CASE							   =    False
					_SCREAMING_SNAKE_CASE							   =    jnp.floataa
					_SCREAMING_SNAKE_CASE							   =    True
					_SCREAMING_SNAKE_CASE							   =    0
					_SCREAMING_SNAKE_CASE							   =    "rgb"
					_SCREAMING_SNAKE_CASE							   =    (16, 32, 96, 256)


					def    SCREAMING_SNAKE_CASE_		(					self	:  Union[str, Any]       ,     __snake_case	:  jax.random.KeyArray       ):





											'''simple docstring'''





											_snake_case:   Dict  					=		(1, self.in_channels, self.sample_size, self.sample_size)
											_snake_case:   Union[str, Any]  					=		jnp.zeros(__snake_case       ,     dtype=jnp.floataa       )
											_snake_case:   Optional[int]  					=		jnp.ones((1,)       ,     dtype=jnp.intaa       )
											_snake_case:   str  					=		jnp.zeros((1, 1, self.cross_attention_dim)       ,     dtype=jnp.floataa       )
											_snake_case:   Dict  					=		(1, 3, self.sample_size * 8, self.sample_size * 8)
											_snake_case:   Dict  					=		jnp.zeros(__snake_case       ,     dtype=jnp.floataa       )

											_snake_case					,						_snake_case:   Any  					=		jax.random.split(__snake_case       )
											_snake_case:   Tuple  					=		{'params': params_rng, 'dropout': dropout_rng}

											return self.init(__snake_case       ,     __snake_case       ,     __snake_case       ,     __snake_case       ,     __snake_case       )["params"]


					def    SCREAMING_SNAKE_CASE_		(					self	:  List[str]       ):





											'''simple docstring'''





											_snake_case:   Any  					=		self.block_out_channels
											_snake_case:   str  					=		block_out_channels[0] * 4

											# If `num_attention_heads` is not defined (which is the case for most models)
											# it will default to `attention_head_dim`. This looks weird upon first reading it and it is.
											# The reason for this behavior is to correct for incorrectly named variables that were introduced
											# when this library was created. The incorrect naming was only discovered much later in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131
											# Changing `attention_head_dim` to `num_attention_heads` for 40,000+ configurations is too backwards breaking
											# which is why we correct for the naming here.
											_snake_case:   Optional[Any]  					=		self.num_attention_heads or self.attention_head_dim

											# input
											_snake_case:   List[Any]  					=		nn.Conv(
											    block_out_channels[0]       ,     kernel_size=(3, 3)       ,     strides=(1, 1)       ,     padding=((1, 1), (1, 1))       ,     dtype=self.dtype       ,     )

											# time
											_snake_case:   int  					=		FlaxTimesteps(
											    block_out_channels[0]       ,     flip_sin_to_cos=self.flip_sin_to_cos       ,     freq_shift=self.config.freq_shift       )
											_snake_case:   List[str]  					=		FlaxTimestepEmbedding(__snake_case       ,     dtype=self.dtype       )

											_snake_case:   List[Any]  					=		FlaxControlNetConditioningEmbedding(
											    conditioning_embedding_channels=block_out_channels[0]       ,     block_out_channels=self.conditioning_embedding_out_channels       ,     )

											_snake_case:   List[Any]  					=		self.only_cross_attention
											if isinstance(__snake_case       ,     __snake_case       ):
																	_snake_case:   int  					=		(only_cross_attention,) * len(self.down_block_types       )

											if isinstance(__snake_case       ,     __snake_case       ):
																	_snake_case:   Tuple  					=		(num_attention_heads,) * len(self.down_block_types       )

											# down
											_snake_case:   List[str]  					=		[]
											_snake_case:   List[str]  					=		[]

											_snake_case:   List[Any]  					=		block_out_channels[0]

											_snake_case:   Optional[Any]  					=		nn.Conv(
											    __snake_case       ,     kernel_size=(1, 1)       ,     padding='VALID'       ,     kernel_init=nn.initializers.zeros_init()       ,     bias_init=nn.initializers.zeros_init()       ,     dtype=self.dtype       ,     )
											controlnet_down_blocks.append(__snake_case       )

											for i, down_block_type in enumerate(self.down_block_types       ):
																	_snake_case:   str  					=		output_channel
																	_snake_case:   Union[str, Any]  					=		block_out_channels[i]
																	_snake_case:   Any  					=		i == len(__snake_case       ) - 1

																	if down_block_type == "CrossAttnDownBlock2D":
																							_snake_case:   int  					=		FlaxCrossAttnDownBlockaD(
																							    in_channels=__snake_case       ,     out_channels=__snake_case       ,     dropout=self.dropout       ,     num_layers=self.layers_per_block       ,     num_attention_heads=num_attention_heads[i]       ,     add_downsample=not is_final_block       ,     use_linear_projection=self.use_linear_projection       ,     only_cross_attention=only_cross_attention[i]       ,     dtype=self.dtype       ,     )
																	else:
																							_snake_case:   Optional[int]  					=		FlaxDownBlockaD(
																							    in_channels=__snake_case       ,     out_channels=__snake_case       ,     dropout=self.dropout       ,     num_layers=self.layers_per_block       ,     add_downsample=not is_final_block       ,     dtype=self.dtype       ,     )

																	down_blocks.append(__snake_case       )

																	for _ in range(self.layers_per_block       ):
																							_snake_case:   Any  					=		nn.Conv(
																							    __snake_case       ,     kernel_size=(1, 1)       ,     padding='VALID'       ,     kernel_init=nn.initializers.zeros_init()       ,     bias_init=nn.initializers.zeros_init()       ,     dtype=self.dtype       ,     )
																							controlnet_down_blocks.append(__snake_case       )

																	if not is_final_block:
																							_snake_case:   str  					=		nn.Conv(
																							    __snake_case       ,     kernel_size=(1, 1)       ,     padding='VALID'       ,     kernel_init=nn.initializers.zeros_init()       ,     bias_init=nn.initializers.zeros_init()       ,     dtype=self.dtype       ,     )
																							controlnet_down_blocks.append(__snake_case       )

											_snake_case:   List[str]  					=		down_blocks
											_snake_case:   Optional[int]  					=		controlnet_down_blocks

											# mid
											_snake_case:   List[Any]  					=		block_out_channels[-1]
											_snake_case:   Tuple  					=		FlaxUNetMidBlockaDCrossAttn(
											    in_channels=__snake_case       ,     dropout=self.dropout       ,     num_attention_heads=num_attention_heads[-1]       ,     use_linear_projection=self.use_linear_projection       ,     dtype=self.dtype       ,     )

											_snake_case:   List[Any]  					=		nn.Conv(
											    __snake_case       ,     kernel_size=(1, 1)       ,     padding='VALID'       ,     kernel_init=nn.initializers.zeros_init()       ,     bias_init=nn.initializers.zeros_init()       ,     dtype=self.dtype       ,     )




					def __call__(					self	:  Any       ,     __snake_case	:  Any       ,     __snake_case	:  List[str]       ,     __snake_case	:  Any       ,     __snake_case	:  Optional[Any]       ,     __snake_case	:  float = 1.0       ,     __snake_case	:  bool = True       ,     __snake_case	:  bool = False       ,     ):





											'''simple docstring'''





											_snake_case:   Dict  					=		self.controlnet_conditioning_channel_order
											if channel_order == "bgr":
																	_snake_case:   int  					=		jnp.flip(__snake_case       ,     axis=1       )

											# 1. time
											if not isinstance(__snake_case       ,     jnp.ndarray       ):
																	_snake_case:   Any  					=		jnp.array([timesteps]       ,     dtype=jnp.intaa       )
											elif isinstance(__snake_case       ,     jnp.ndarray       ) and len(timesteps.shape       ) == 0:
																	_snake_case:   Any  					=		timesteps.astype(dtype=jnp.floataa       )
																	_snake_case:   Tuple  					=		jnp.expand_dims(__snake_case       ,     0       )

											_snake_case:   Optional[int]  					=		self.time_proj(__snake_case       )
											_snake_case:   Union[str, Any]  					=		self.time_embedding(__snake_case       )

											# 2. pre-process
											_snake_case:   str  					=		jnp.transpose(__snake_case       ,     (0, 2, 3, 1)       )
											_snake_case:   Dict  					=		self.conv_in(__snake_case       )

											_snake_case:   Union[str, Any]  					=		jnp.transpose(__snake_case       ,     (0, 2, 3, 1)       )
											_snake_case:   List[Any]  					=		self.controlnet_cond_embedding(__snake_case       )
											sample += controlnet_cond

											# 3. down
											_snake_case:   Any  					=		(sample,)
											for down_block in self.down_blocks:
																	if isinstance(__snake_case       ,     __snake_case       ):
																							_snake_case					,						_snake_case:   Dict  					=		down_block(__snake_case       ,     __snake_case       ,     __snake_case       ,     deterministic=not train       )
																	else:
																							_snake_case					,						_snake_case:   Union[str, Any]  					=		down_block(__snake_case       ,     __snake_case       ,     deterministic=not train       )
																	down_block_res_samples += res_samples

											# 4. mid
											_snake_case:   List[str]  					=		self.mid_block(__snake_case       ,     __snake_case       ,     __snake_case       ,     deterministic=not train       )

											# 5. contronet blocks
											_snake_case:   Dict  					=		()
											for down_block_res_sample, controlnet_block in zip(__snake_case       ,     self.controlnet_down_blocks       ):
																	_snake_case:   List[str]  					=		controlnet_block(__snake_case       )
																	controlnet_down_block_res_samples += (down_block_res_sample,)

											_snake_case:   Optional[int]  					=		controlnet_down_block_res_samples

											_snake_case:   Optional[Any]  					=		self.controlnet_mid_block(__snake_case       )

											# 6. scaling
											_snake_case:   str  					=		[sample * conditioning_scale for sample in down_block_res_samples]
											mid_block_res_sample *= conditioning_scale

											if not return_dict:
																	return (down_block_res_samples, mid_block_res_sample)

											return FlaxControlNetOutput(
											    down_block_res_samples=__snake_case       ,     mid_block_res_sample=__snake_case       )












'''simple docstring'''






from __future__ import annotations

import sys
from collections import deque
from typing import Generic, TypeVar

A   : str             =			TypeVar('T')







class 			lowerCamelCase					(						Generic[T]  ):
					_SCREAMING_SNAKE_CASE							   =    42  # Cache store of keys
					_SCREAMING_SNAKE_CASE							   =    42  # References of the keys in cache
					_SCREAMING_SNAKE_CASE							   =    10  # Maximum capacity of cache


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





											'''simple docstring'''





											_snake_case:   Dict  					=		deque()
											_snake_case:   Union[str, Any]  					=		set()
											if not n:
																	_snake_case:   Optional[int]  					=		sys.maxsize
											elif n < 0:
																	raise ValueError('n should be an integer greater than 0.'       )
											else:
																	_snake_case:   Tuple  					=		n


					def    SCREAMING_SNAKE_CASE_		(					self	:  Any       ,     __snake_case	:  T       ):





											'''simple docstring'''





											if x not in self.key_reference:
																	if len(self.dq_store       ) == LRUCache._MAX_CAPACITY:
																							_snake_case:   int  					=		self.dq_store.pop()
																							self.key_reference.remove(__snake_case       )
											else:
																	self.dq_store.remove(__snake_case       )

											self.dq_store.appendleft(__snake_case       )
											self.key_reference.add(__snake_case       )


					def    SCREAMING_SNAKE_CASE_		(					self	:  int       ):





											'''simple docstring'''





											for k in self.dq_store:
																	print(__snake_case       )




					def __repr__(					self	:  List[Any]       ):





											'''simple docstring'''





											return f'''LRUCache({self._MAX_CAPACITY}) => {list(self.dq_store       )}'''


if __name__ == "__main__":
		import doctest

		doctest.testmod()

		A   : LRUCache[str | int]             =			LRUCache(4)
		lru_cache.refer('A')
		lru_cache.refer(2)
		lru_cache.refer(3)
		lru_cache.refer('A')
		lru_cache.refer(4)
		lru_cache.refer(5)
		lru_cache.display()

		print(lru_cache)
		assert str(lru_cache) == "LRUCache(4) => [5, 4, 'A', 3]"





