



"""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     UpperCAmelCase    (	__SCREAMING_SNAKE_CASE  ):
   A__		:	jnp.ndarray
   A__		:	jnp.ndarray



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



   def    __UpperCAmelCase    (       self    :	List[str]				):



     """simple docstring"""
     _snake_case   	=      nn.Conv(
         self.block_out_channels[0]							,       kernel_size=(3, 3)							,       padding=((1, 1), (1, 1))							,       dtype=self.dtype							,       )

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

     _snake_case   	=      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    :	int							,       __lowerCamelCase    :	int				):



     """simple docstring"""
     _snake_case   	=      self.conv_in(__lowerCamelCase				)
     _snake_case   	=      nn.silu(__lowerCamelCase				)

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

     _snake_case   	=      self.conv_out(__lowerCamelCase				)

     return embedding



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



   def    __UpperCAmelCase    (       self    :	Any							,       __lowerCamelCase    :	jax.random.KeyArray				):



     """simple docstring"""
     # init input tensors
     _snake_case   	=      (1, self.in_channels, self.sample_size, self.sample_size)
     _snake_case   	=      jnp.zeros(__lowerCamelCase							,       dtype=jnp.floataa				)
     _snake_case   	=      jnp.ones((1,)							,       dtype=jnp.intaa				)
     _snake_case   	=      jnp.zeros((1, 1, self.cross_attention_dim)							,       dtype=jnp.floataa				)
     _snake_case   	=      (1, 3, self.sample_size * 8, self.sample_size * 8)
     _snake_case   	=      jnp.zeros(__lowerCamelCase							,       dtype=jnp.floataa				)

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

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



   def    __UpperCAmelCase    (       self    :	int				):



     """simple docstring"""
     _snake_case   	=      self.block_out_channels
     _snake_case   	=      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   	=      self.num_attention_heads or self.attention_head_dim

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

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

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

     _snake_case   	=      self.only_cross_attention
     if isinstance(__lowerCamelCase							,       __lowerCamelCase				):
       _snake_case   	=      (only_cross_attention,) * len(self.down_block_types				)

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

     # down
     _snake_case   	=      []
     _snake_case   	=      []

     _snake_case   	=      block_out_channels[0]

     _snake_case   	=      nn.Conv(
         __lowerCamelCase							,       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(__lowerCamelCase				)

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

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

       down_blocks.append(__lowerCamelCase				)

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

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

     _snake_case   	=      down_blocks
     _snake_case   	=      controlnet_down_blocks

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

     _snake_case   	=      nn.Conv(
         __lowerCamelCase							,       kernel_size=(1, 1)							,       padding='''VALID'''							,       kernel_init=nn.initializers.zeros_init()							,       bias_init=nn.initializers.zeros_init()							,       dtype=self.dtype							,       )





   def __call__(       self    :	List[Any]							,       __lowerCamelCase    :	int							,       __lowerCamelCase    :	Any							,       __lowerCamelCase    :	Union[str, Any]							,       __lowerCamelCase    :	Tuple							,       __lowerCamelCase    :	float = 1.0							,       __lowerCamelCase    :	bool = True							,       __lowerCamelCase    :	bool = False							,       ):



     """simple docstring"""
     _snake_case   	=      self.controlnet_conditioning_channel_order
     if channel_order == "bgr":
       _snake_case   	=      jnp.flip(__lowerCamelCase							,       axis=1				)

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

     _snake_case   	=      self.time_proj(__lowerCamelCase				)
     _snake_case   	=      self.time_embedding(__lowerCamelCase				)

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

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

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

     # 4. mid
     _snake_case   	=      self.mid_block(__lowerCamelCase							,       __lowerCamelCase							,       __lowerCamelCase							,       deterministic=not train				)

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

     _snake_case   	=      controlnet_down_block_res_samples

     _snake_case   	=      self.controlnet_mid_block(__lowerCamelCase				)

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










"""simple docstring"""






from ..utils import DummyObject, requires_backends



class     UpperCAmelCase    (	metaclass=__SCREAMING_SNAKE_CASE  ):
   A__		:	Optional[int]           =  ['''torch''', '''scipy''']



   def __init__(       self    :	Any							,       *__lowerCamelCase    :	List[Any]							,       **__lowerCamelCase    :	Any				):



     """simple docstring"""
     requires_backends(self							,       ['''torch''', '''scipy''']				)



   @classmethod
   def    __UpperCAmelCase    (       cls    :	Dict							,       *__lowerCamelCase    :	List[str]							,       **__lowerCamelCase    :	Tuple				):



     """simple docstring"""
     requires_backends(cls							,       ['''torch''', '''scipy''']				)





   @classmethod
   def    __UpperCAmelCase    (       cls    :	int							,       *__lowerCamelCase    :	Union[str, Any]							,       **__lowerCamelCase    :	Tuple				):



     """simple docstring"""
     requires_backends(cls							,       ['''torch''', '''scipy''']				)





