





from __future__ import annotations

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

UpperCamelCase_												=	TypeVar('T')
UpperCamelCase_												=	TypeVar('U')

class     snake_case_		(   Generic[T, U]							):






 '''simple docstring'''





 def __init__(  self,		A_,		A_		)	->       str:
    UpperCAmelCase__						=key
    UpperCAmelCase__						=val
    UpperCAmelCase__						=None
    UpperCAmelCase__						=None






 def __repr__(  self		)	->       str:
    return (
        f"""Node: key: {self.key}, val: {self.val}, """
        f"""has next: {bool(self.next		)}, has prev: {bool(self.prev		)}"""
    )

class     snake_case_		(   Generic[T, U]							):






 '''simple docstring'''





 def __init__(  self		)	->       None:
    UpperCAmelCase__						=DoubleLinkedListNode(A_,		A_		)
    UpperCAmelCase__						=DoubleLinkedListNode(A_,		A_		)
    UpperCAmelCase__			,							UpperCAmelCase__						=self.rear, self.head

 def __repr__(  self		)	->       str:
    UpperCAmelCase__						=["DoubleLinkedList"]
    UpperCAmelCase__						=self.head
    while node.next is not None:
       rep.append(str(A_		)		)
       UpperCAmelCase__						=node.next
    rep.append(str(self.rear		)		)
    return ",\n    ".join(A_		)

 def  __UpperCAmelCase						(  self,		A_		)	->       None:
    UpperCAmelCase__						=self.rear.prev

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

    UpperCAmelCase__						=node
    UpperCAmelCase__						=previous
    UpperCAmelCase__						=node
    UpperCAmelCase__						=self.rear






 def  __UpperCAmelCase						(  self,		A_		)	->       DoubleLinkedListNode[T, U] | None:

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

    UpperCAmelCase__						=node.next
    UpperCAmelCase__						=node.prev
    UpperCAmelCase__						=None
    UpperCAmelCase__						=None
    return node

class     snake_case_		(   Generic[T, U]							):






 '''simple docstring'''





 __UpperCamelCase       =      {}

 def __init__(  self,		A_		)	->       Tuple:
    UpperCAmelCase__						=DoubleLinkedList()
    UpperCAmelCase__						=capacity
    UpperCAmelCase__						=0
    UpperCAmelCase__						=0
    UpperCAmelCase__						=0
    UpperCAmelCase__						={}

 def __repr__(  self		)	->       str:
    return (
        f"""CacheInfo(hits={self.hits}, misses={self.miss}, """
        f"""capacity={self.capacity}, current size={self.num_keys})"""
    )

 def __contains__(  self,		A_		)	->       bool:
    return key in self.cache

 def  __UpperCAmelCase						(  self,		A_		)	->       U | None:
    # Note: pythonic interface would throw KeyError rather than return None

    if key in self.cache:
       self.hits += 1
       UpperCAmelCase__						=self.cache[key]
       UpperCAmelCase__						=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  __UpperCAmelCase						(  self,		A_,		A_		)	->       None:

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

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






 @classmethod
 def  __UpperCAmelCase						(  cls,		A_ = 128		)	->       Callable[[Callable[[T], U]], Callable[..., U]]:

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

          UpperCAmelCase__						=cls.decorator_function_to_instance_map[func].get(args[0]		)
          if result is None:
             UpperCAmelCase__						=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()










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     snake_case_		(   a							):






 '''simple docstring'''





 __UpperCamelCase       =      42
 __UpperCamelCase       =      42

class     snake_case_		(   nn.Module							):






 '''simple docstring'''





 __UpperCamelCase       =      42
 __UpperCamelCase       =      (1_6, 3_2, 9_6, 2_5_6)
 __UpperCamelCase       =      jnp.floataa

 def  __UpperCAmelCase						(  self		)	->       Any:
    UpperCAmelCase__						=nn.Conv(
        self.block_out_channels[0],		kernel_size=(3, 3),		padding=((1, 1), (1, 1)),		dtype=self.dtype,		)

    UpperCAmelCase__						=[]
    for i in range(len(self.block_out_channels		) - 1		):
       UpperCAmelCase__						=self.block_out_channels[i]
       UpperCAmelCase__						=self.block_out_channels[i + 1]
       UpperCAmelCase__						=nn.Conv(
           A_,		kernel_size=(3, 3),		padding=((1, 1), (1, 1)),		dtype=self.dtype,		)
       blocks.append(A_		)
       UpperCAmelCase__						=nn.Conv(
           A_,		kernel_size=(3, 3),		strides=(2, 2),		padding=((1, 1), (1, 1)),		dtype=self.dtype,		)
       blocks.append(A_		)
    UpperCAmelCase__						=blocks

    UpperCAmelCase__						=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,		A_		)	->       Any:
    UpperCAmelCase__						=self.conv_in(A_		)
    UpperCAmelCase__						=nn.silu(A_		)

    for block in self.blocks:
       UpperCAmelCase__						=block(A_		)
       UpperCAmelCase__						=nn.silu(A_		)

    UpperCAmelCase__						=self.conv_out(A_		)

    return embedding

@flax_register_to_config
class     snake_case_		(   nn.Module,		a,		a							):






 '''simple docstring'''





 __UpperCamelCase       =      3_2
 __UpperCamelCase       =      4
 __UpperCamelCase       =      (
     "CrossAttnDownBlock2D",
     "CrossAttnDownBlock2D",
     "CrossAttnDownBlock2D",
     "DownBlock2D",
 )
 __UpperCamelCase       =      False
 __UpperCamelCase       =      (3_2_0, 6_4_0, 1_2_8_0, 1_2_8_0)
 __UpperCamelCase       =      2
 __UpperCamelCase       =      8
 __UpperCamelCase       =      None
 __UpperCamelCase       =      1_2_8_0
 __UpperCamelCase       =      0.0
 __UpperCamelCase       =      False
 __UpperCamelCase       =      jnp.floataa
 __UpperCamelCase       =      True
 __UpperCamelCase       =      0
 __UpperCamelCase       =      "rgb"
 __UpperCamelCase       =      (1_6, 3_2, 9_6, 2_5_6)

 def  __UpperCAmelCase						(  self,		A_		)	->       FrozenDict:
    # init input tensors
    UpperCAmelCase__						=(1, self.in_channels, self.sample_size, self.sample_size)
    UpperCAmelCase__						=jnp.zeros(A_,		dtype=jnp.floataa		)
    UpperCAmelCase__						=jnp.ones((1,),		dtype=jnp.intaa		)
    UpperCAmelCase__						=jnp.zeros((1, 1, self.cross_attention_dim),		dtype=jnp.floataa		)
    UpperCAmelCase__						=(1, 3, self.sample_size * 8, self.sample_size * 8)
    UpperCAmelCase__						=jnp.zeros(A_,		dtype=jnp.floataa		)

    UpperCAmelCase__			,							UpperCAmelCase__						=jax.random.split(A_		)
    UpperCAmelCase__						={"params": params_rng, "dropout": dropout_rng}

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

 def  __UpperCAmelCase						(  self		)	->       int:
    UpperCAmelCase__						=self.block_out_channels
    UpperCAmelCase__						=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.
    UpperCAmelCase__						=self.num_attention_heads or self.attention_head_dim

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

    # time
    UpperCAmelCase__						=FlaxTimesteps(
        block_out_channels[0],		flip_sin_to_cos=self.flip_sin_to_cos,		freq_shift=self.config.freq_shift		)
    UpperCAmelCase__						=FlaxTimestepEmbedding(A_,		dtype=self.dtype		)

    UpperCAmelCase__						=FlaxControlNetConditioningEmbedding(
        conditioning_embedding_channels=block_out_channels[0],		block_out_channels=self.conditioning_embedding_out_channels,		)

    UpperCAmelCase__						=self.only_cross_attention
    if isinstance(A_,		A_		):
       UpperCAmelCase__						=(only_cross_attention,) * len(self.down_block_types		)

    if isinstance(A_,		A_		):
       UpperCAmelCase__						=(num_attention_heads,) * len(self.down_block_types		)

    # down
    UpperCAmelCase__						=[]
    UpperCAmelCase__						=[]

    UpperCAmelCase__						=block_out_channels[0]

    UpperCAmelCase__						=nn.Conv(
        A_,		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(A_		)

    for i, down_block_type in enumerate(self.down_block_types		):
       UpperCAmelCase__						=output_channel
       UpperCAmelCase__						=block_out_channels[i]
       UpperCAmelCase__						=i == len(A_		) - 1

       if down_block_type == "CrossAttnDownBlock2D":
          UpperCAmelCase__						=FlaxCrossAttnDownBlockaD(
              in_channels=A_,		out_channels=A_,		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:
          UpperCAmelCase__						=FlaxDownBlockaD(
              in_channels=A_,		out_channels=A_,		dropout=self.dropout,		num_layers=self.layers_per_block,		add_downsample=not is_final_block,		dtype=self.dtype,		)

       down_blocks.append(A_		)

       for _ in range(self.layers_per_block		):
          UpperCAmelCase__						=nn.Conv(
              A_,		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(A_		)

       if not is_final_block:
          UpperCAmelCase__						=nn.Conv(
              A_,		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(A_		)

    UpperCAmelCase__						=down_blocks
    UpperCAmelCase__						=controlnet_down_blocks

    # mid
    UpperCAmelCase__						=block_out_channels[-1]
    UpperCAmelCase__						=FlaxUNetMidBlockaDCrossAttn(
        in_channels=A_,		dropout=self.dropout,		num_attention_heads=num_attention_heads[-1],		use_linear_projection=self.use_linear_projection,		dtype=self.dtype,		)

    UpperCAmelCase__						=nn.Conv(
        A_,		kernel_size=(1, 1),		padding="VALID",		kernel_init=nn.initializers.zeros_init(),		bias_init=nn.initializers.zeros_init(),		dtype=self.dtype,		)






 def __call__(  self,		A_,		A_,		A_,		A_,		A_ = 1.0,		A_ = True,		A_ = False,		)	->       Union[FlaxControlNetOutput, Tuple]:
    UpperCAmelCase__						=self.controlnet_conditioning_channel_order
    if channel_order == "bgr":
       UpperCAmelCase__						=jnp.flip(A_,		axis=1		)

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

    UpperCAmelCase__						=self.time_proj(A_		)
    UpperCAmelCase__						=self.time_embedding(A_		)

    # 2. pre-process
    UpperCAmelCase__						=jnp.transpose(A_,		(0, 2, 3, 1)		)
    UpperCAmelCase__						=self.conv_in(A_		)

    UpperCAmelCase__						=jnp.transpose(A_,		(0, 2, 3, 1)		)
    UpperCAmelCase__						=self.controlnet_cond_embedding(A_		)
    sample += controlnet_cond

    # 3. down
    UpperCAmelCase__						=(sample,)
    for down_block in self.down_blocks:
       if isinstance(A_,		A_		):
          UpperCAmelCase__			,							UpperCAmelCase__						=down_block(A_,		A_,		A_,		deterministic=not train		)
       else:
          UpperCAmelCase__			,							UpperCAmelCase__						=down_block(A_,		A_,		deterministic=not train		)
       down_block_res_samples += res_samples

    # 4. mid
    UpperCAmelCase__						=self.mid_block(A_,		A_,		A_,		deterministic=not train		)

    # 5. contronet blocks
    UpperCAmelCase__						=()
    for down_block_res_sample, controlnet_block in zip(A_,		self.controlnet_down_blocks		):
       UpperCAmelCase__						=controlnet_block(A_		)
       controlnet_down_block_res_samples += (down_block_res_sample,)

    UpperCAmelCase__						=controlnet_down_block_res_samples

    UpperCAmelCase__						=self.controlnet_mid_block(A_		)

    # 6. scaling
    UpperCAmelCase__						=[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=A_,		mid_block_res_sample=A_		)



