



"""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       SCREAMING_SNAKE_CASE__  (					UpperCAmelCase__ ):
   __lowerCAmelCase  :   jnp.ndarray
   __lowerCAmelCase  :   jnp.ndarray



class       SCREAMING_SNAKE_CASE__  (					nn.Module ):
   __lowerCAmelCase  :   int
   __lowerCAmelCase  :   Tuple[int]						=						(16, 32, 96, 256)
   __lowerCAmelCase  :   jnp.dtype						=						jnp.floataa

   def 	SCREAMING_SNAKE_CASE     (   self						)  ->      str:





    '''simple docstring'''




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

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

    UpperCAmelCase	:					Optional[int]	=	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      ,     _SCREAMING_SNAKE_CASE						)  ->      Dict:





    '''simple docstring'''




    UpperCAmelCase	:					Dict	=	self.conv_in(_SCREAMING_SNAKE_CASE						)
    UpperCAmelCase	:					Dict	=	nn.silu(_SCREAMING_SNAKE_CASE						)

    for block in self.blocks:
     UpperCAmelCase	:					Dict	=	block(_SCREAMING_SNAKE_CASE						)
     UpperCAmelCase	:					int	=	nn.silu(_SCREAMING_SNAKE_CASE						)

    UpperCAmelCase	:					List[Any]	=	self.conv_out(_SCREAMING_SNAKE_CASE						)

    return embedding



@flax_register_to_config
class       SCREAMING_SNAKE_CASE__  (					nn.Module      ,			UpperCAmelCase__      ,			UpperCAmelCase__ ):
   __lowerCAmelCase  :   int						=						32
   __lowerCAmelCase  :   int						=						4
   __lowerCAmelCase  :   Tuple[str]						=						(
       "CrossAttnDownBlock2D",
       "CrossAttnDownBlock2D",
       "CrossAttnDownBlock2D",
       "DownBlock2D",
   )
   __lowerCAmelCase  :   Union[bool, Tuple[bool]]						=						False
   __lowerCAmelCase  :   Tuple[int]						=						(320, 640, 1_280, 1_280)
   __lowerCAmelCase  :   int						=						2
   __lowerCAmelCase  :   Union[int, Tuple[int]]						=						8
   __lowerCAmelCase  :   Optional[Union[int, Tuple[int]]]						=						None
   __lowerCAmelCase  :   int						=						1_280
   __lowerCAmelCase  :   float						=						0.0
   __lowerCAmelCase  :   bool						=						False
   __lowerCAmelCase  :   jnp.dtype						=						jnp.floataa
   __lowerCAmelCase  :   bool						=						True
   __lowerCAmelCase  :   int						=						0
   __lowerCAmelCase  :   str						=						"rgb"
   __lowerCAmelCase  :   Tuple[int]						=						(16, 32, 96, 256)

   def 	SCREAMING_SNAKE_CASE     (   self      ,     _SCREAMING_SNAKE_CASE						)  ->      FrozenDict:





    '''simple docstring'''




    UpperCAmelCase	:					Optional[int]	=	(1, self.in_channels, self.sample_size, self.sample_size)
    UpperCAmelCase	:					Any	=	jnp.zeros(_SCREAMING_SNAKE_CASE      ,     dtype=jnp.floataa						)
    UpperCAmelCase	:					List[Any]	=	jnp.ones((1,)      ,     dtype=jnp.intaa						)
    UpperCAmelCase	:					Dict	=	jnp.zeros((1, 1, self.cross_attention_dim)      ,     dtype=jnp.floataa						)
    UpperCAmelCase	:					Tuple	=	(1, 3, self.sample_size * 8, self.sample_size * 8)
    UpperCAmelCase	:					str	=	jnp.zeros(_SCREAMING_SNAKE_CASE      ,     dtype=jnp.floataa						)

    UpperCAmelCase	,		UpperCAmelCase	:					List[Any]	=	jax.random.split(_SCREAMING_SNAKE_CASE						)
    UpperCAmelCase	:					Dict	=	{"""params""": params_rng, """dropout""": dropout_rng}

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

   def 	SCREAMING_SNAKE_CASE     (   self						)  ->      Optional[int]:





    '''simple docstring'''




    UpperCAmelCase	:					Optional[Any]	=	self.block_out_channels
    UpperCAmelCase	:					List[Any]	=	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	:					str	=	self.num_attention_heads or self.attention_head_dim

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

    # time
    UpperCAmelCase	:					Union[str, Any]	=	FlaxTimesteps(
        block_out_channels[0]      ,     flip_sin_to_cos=self.flip_sin_to_cos      ,     freq_shift=self.config.freq_shift						)
    UpperCAmelCase	:					List[str]	=	FlaxTimestepEmbedding(_SCREAMING_SNAKE_CASE      ,     dtype=self.dtype						)

    UpperCAmelCase	:					Optional[Any]	=	FlaxControlNetConditioningEmbedding(
        conditioning_embedding_channels=block_out_channels[0]      ,     block_out_channels=self.conditioning_embedding_out_channels      ,     )

    UpperCAmelCase	:					Tuple	=	self.only_cross_attention
    if isinstance(_SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE						):
     UpperCAmelCase	:					Tuple	=	(only_cross_attention,) * len(self.down_block_types						)

    if isinstance(_SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE						):
     UpperCAmelCase	:					Tuple	=	(num_attention_heads,) * len(self.down_block_types						)

    # down
    UpperCAmelCase	:					Optional[int]	=	[]
    UpperCAmelCase	:					Optional[int]	=	[]

    UpperCAmelCase	:					List[Any]	=	block_out_channels[0]

    UpperCAmelCase	:					Optional[Any]	=	nn.Conv(
        _SCREAMING_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(_SCREAMING_SNAKE_CASE						)

    for i, down_block_type in enumerate(self.down_block_types						):
     UpperCAmelCase	:					Optional[int]	=	output_channel
     UpperCAmelCase	:					Dict	=	block_out_channels[i]
     UpperCAmelCase	:					Optional[Any]	=	i == len(_SCREAMING_SNAKE_CASE						) - 1

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

     down_blocks.append(_SCREAMING_SNAKE_CASE						)

     for _ in range(self.layers_per_block						):
      UpperCAmelCase	:					Tuple	=	nn.Conv(
          _SCREAMING_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(_SCREAMING_SNAKE_CASE						)

     if not is_final_block:
      UpperCAmelCase	:					Tuple	=	nn.Conv(
          _SCREAMING_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(_SCREAMING_SNAKE_CASE						)

    UpperCAmelCase	:					List[Any]	=	down_blocks
    UpperCAmelCase	:					List[Any]	=	controlnet_down_blocks

    # mid
    UpperCAmelCase	:					Optional[int]	=	block_out_channels[-1]
    UpperCAmelCase	:					List[str]	=	FlaxUNetMidBlockaDCrossAttn(
        in_channels=_SCREAMING_SNAKE_CASE      ,     dropout=self.dropout      ,     num_attention_heads=num_attention_heads[-1]      ,     use_linear_projection=self.use_linear_projection      ,     dtype=self.dtype      ,     )

    UpperCAmelCase	:					Optional[int]	=	nn.Conv(
        _SCREAMING_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      ,     _SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE = 1.0      ,     _SCREAMING_SNAKE_CASE = True      ,     _SCREAMING_SNAKE_CASE = False      ,     )  ->      Union[FlaxControlNetOutput, Tuple]:





    '''simple docstring'''




    UpperCAmelCase	:					List[str]	=	self.controlnet_conditioning_channel_order
    if channel_order == "bgr":
     UpperCAmelCase	:					str	=	jnp.flip(_SCREAMING_SNAKE_CASE      ,     axis=1						)

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

    UpperCAmelCase	:					Union[str, Any]	=	self.time_proj(_SCREAMING_SNAKE_CASE						)
    UpperCAmelCase	:					Any	=	self.time_embedding(_SCREAMING_SNAKE_CASE						)

    # 2. pre-process
    UpperCAmelCase	:					Dict	=	jnp.transpose(_SCREAMING_SNAKE_CASE      ,     (0, 2, 3, 1)						)
    UpperCAmelCase	:					List[str]	=	self.conv_in(_SCREAMING_SNAKE_CASE						)

    UpperCAmelCase	:					Any	=	jnp.transpose(_SCREAMING_SNAKE_CASE      ,     (0, 2, 3, 1)						)
    UpperCAmelCase	:					Optional[Any]	=	self.controlnet_cond_embedding(_SCREAMING_SNAKE_CASE						)
    sample += controlnet_cond

    # 3. down
    UpperCAmelCase	:					List[Any]	=	(sample,)
    for down_block in self.down_blocks:
     if isinstance(_SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE						):
      UpperCAmelCase	,		UpperCAmelCase	:					Optional[int]	=	down_block(_SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE      ,     deterministic=not train						)
     else:
      UpperCAmelCase	,		UpperCAmelCase	:					str	=	down_block(_SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE      ,     deterministic=not train						)
     down_block_res_samples += res_samples

    # 4. mid
    UpperCAmelCase	:					Dict	=	self.mid_block(_SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE      ,     deterministic=not train						)

    # 5. contronet blocks
    UpperCAmelCase	:					List[Any]	=	()
    for down_block_res_sample, controlnet_block in zip(_SCREAMING_SNAKE_CASE      ,     self.controlnet_down_blocks						):
     UpperCAmelCase	:					List[str]	=	controlnet_block(_SCREAMING_SNAKE_CASE						)
     controlnet_down_block_res_samples += (down_block_res_sample,)

    UpperCAmelCase	:					List[Any]	=	controlnet_down_block_res_samples

    UpperCAmelCase	:					Optional[Any]	=	self.controlnet_mid_block(_SCREAMING_SNAKE_CASE						)

    # 6. scaling
    UpperCAmelCase	:					Any	=	[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=_SCREAMING_SNAKE_CASE      ,     mid_block_res_sample=_SCREAMING_SNAKE_CASE						)







from __future__ import annotations




def 		lowerCAmelCase_    (  __UpperCAmelCase:			int		,   __UpperCAmelCase:			int      )      ->	list[list[int]]:
    UpperCamelCase__							:				list[list[int]]          =      []
    create_all_state(1		,   __UpperCAmelCase		,   __UpperCAmelCase		,   []		,   __UpperCAmelCase      )
    return result




def 		lowerCAmelCase_    (  __UpperCAmelCase:			int		,   __UpperCAmelCase:			int		,   __UpperCAmelCase:			int		,   __UpperCAmelCase:			list[int]		,   __UpperCAmelCase:			list[list[int]]		,   )      ->	None:
    if level == 0:
        total_list.append(current_list[:]      )
        return

    for i in range(__UpperCAmelCase		,   total_number - level + 2      ):
        current_list.append(__UpperCAmelCase      )
        create_all_state(i + 1		,   __UpperCAmelCase		,   level - 1		,   __UpperCAmelCase		,   __UpperCAmelCase      )
        current_list.pop()




def 		lowerCAmelCase_    (  __UpperCAmelCase:			list[list[int]]      )      ->	None:
    for i in total_list:
        print(*__UpperCAmelCase      )


if __name__ == "__main__":
  UpperCAmelCase_     							= 4
  UpperCAmelCase_     							= 2
  UpperCAmelCase_     							= generate_all_combinations(n, k)
  print_all_state(total_list)

