'''
    #################################################################################################################################################################################################
    # This code library is an adaptation of the original Transformers and was designed, developed and programmed by Sapiens Technology®.                                                            #
    # Any alteration and/or disclosure of this code without prior authorization is strictly prohibited and is subject to legal action that will be forwarded by the Sapiens Technology® legal team. #
    # This set of algorithms aims to download, train, fine-tune and/or infer large language models from various sources and slopes.                                                                 #
    #################################################################################################################################################################################################
'''
from ..unets.unet_2d_blocks_flax import FlaxCrossAttnDownBlock2D, FlaxDownBlock2D, FlaxUNetMidBlock2DCrossAttn
from ...configuration_utils import ConfigMixin, flax_register_to_config
from ..embeddings_flax import FlaxTimestepEmbedding, FlaxTimesteps
from ..modeling_flax_utils import FlaxModelMixin
from flax.core.frozen_dict import FrozenDict
from typing import Optional, Tuple, Union
from ...utils import BaseOutput
import flax.linen as nn
import jax.numpy as jnp
import flax
import jax
@flax.struct.dataclass
class FlaxControlNetOutput(BaseOutput):
    """Args:"""
    down_block_res_samples: jnp.ndarray
    mid_block_res_sample: jnp.ndarray
class FlaxControlNetConditioningEmbedding(nn.Module):
    conditioning_embedding_channels: int
    block_out_channels: Tuple[int, ...] = (16, 32, 96, 256)
    dtype: jnp.dtype = jnp.float32
    def setup(self) -> None:
        self.conv_in = nn.Conv(self.block_out_channels[0], kernel_size=(3, 3), padding=((1, 1), (1, 1)), dtype=self.dtype)
        blocks = []
        for i in range(len(self.block_out_channels) - 1):
            channel_in = self.block_out_channels[i]
            channel_out = self.block_out_channels[i + 1]
            conv1 = nn.Conv(channel_in, kernel_size=(3, 3), padding=((1, 1), (1, 1)), dtype=self.dtype)
            blocks.append(conv1)
            conv2 = nn.Conv(channel_out, kernel_size=(3, 3), strides=(2, 2), padding=((1, 1), (1, 1)), dtype=self.dtype)
            blocks.append(conv2)
        self.blocks = blocks
        self.conv_out = 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, conditioning: jnp.ndarray) -> jnp.ndarray:
        embedding = self.conv_in(conditioning)
        embedding = nn.silu(embedding)
        for block in self.blocks:
            embedding = block(embedding)
            embedding = nn.silu(embedding)
        embedding = self.conv_out(embedding)
        return embedding
@flax_register_to_config
class FlaxControlNetModel(nn.Module, FlaxModelMixin, ConfigMixin):
    sample_size: int = 32
    in_channels: int = 4
    down_block_types: Tuple[str, ...] = ('CrossAttnDownBlock2D', 'CrossAttnDownBlock2D', 'CrossAttnDownBlock2D', 'DownBlock2D')
    only_cross_attention: Union[bool, Tuple[bool, ...]] = False
    block_out_channels: Tuple[int, ...] = (320, 640, 1280, 1280)
    layers_per_block: int = 2
    attention_head_dim: Union[int, Tuple[int, ...]] = 8
    num_attention_heads: Optional[Union[int, Tuple[int, ...]]] = None
    cross_attention_dim: int = 1280
    dropout: float = 0.0
    use_linear_projection: bool = False
    dtype: jnp.dtype = jnp.float32
    flip_sin_to_cos: bool = True
    freq_shift: int = 0
    controlnet_conditioning_channel_order: str = 'rgb'
    conditioning_embedding_out_channels: Tuple[int, ...] = (16, 32, 96, 256)
    def init_weights(self, rng: jax.Array) -> FrozenDict:
        sample_shape = (1, self.in_channels, self.sample_size, self.sample_size)
        sample = jnp.zeros(sample_shape, dtype=jnp.float32)
        timesteps = jnp.ones((1,), dtype=jnp.int32)
        encoder_hidden_states = jnp.zeros((1, 1, self.cross_attention_dim), dtype=jnp.float32)
        controlnet_cond_shape = (1, 3, self.sample_size * 8, self.sample_size * 8)
        controlnet_cond = jnp.zeros(controlnet_cond_shape, dtype=jnp.float32)
        params_rng, dropout_rng = jax.random.split(rng)
        rngs = {'params': params_rng, 'dropout': dropout_rng}
        return self.init(rngs, sample, timesteps, encoder_hidden_states, controlnet_cond)['params']
    def setup(self) -> None:
        block_out_channels = self.block_out_channels
        time_embed_dim = block_out_channels[0] * 4
        num_attention_heads = self.num_attention_heads or self.attention_head_dim
        self.conv_in = nn.Conv(block_out_channels[0], kernel_size=(3, 3), strides=(1, 1), padding=((1, 1), (1, 1)), dtype=self.dtype)
        self.time_proj = FlaxTimesteps(block_out_channels[0], flip_sin_to_cos=self.flip_sin_to_cos, freq_shift=self.config.freq_shift)
        self.time_embedding = FlaxTimestepEmbedding(time_embed_dim, dtype=self.dtype)
        self.controlnet_cond_embedding = FlaxControlNetConditioningEmbedding(conditioning_embedding_channels=block_out_channels[0], block_out_channels=self.conditioning_embedding_out_channels)
        only_cross_attention = self.only_cross_attention
        if isinstance(only_cross_attention, bool): only_cross_attention = (only_cross_attention,) * len(self.down_block_types)
        if isinstance(num_attention_heads, int): num_attention_heads = (num_attention_heads,) * len(self.down_block_types)
        down_blocks = []
        controlnet_down_blocks = []
        output_channel = block_out_channels[0]
        controlnet_block = nn.Conv(output_channel, 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(controlnet_block)
        for i, down_block_type in enumerate(self.down_block_types):
            input_channel = output_channel
            output_channel = block_out_channels[i]
            is_final_block = i == len(block_out_channels) - 1
            if down_block_type == 'CrossAttnDownBlock2D': down_block = FlaxCrossAttnDownBlock2D(in_channels=input_channel, out_channels=output_channel, 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: down_block = FlaxDownBlock2D(in_channels=input_channel, out_channels=output_channel, dropout=self.dropout, num_layers=self.layers_per_block, add_downsample=not is_final_block, dtype=self.dtype)
            down_blocks.append(down_block)
            for _ in range(self.layers_per_block):
                controlnet_block = nn.Conv(output_channel, 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(controlnet_block)
            if not is_final_block:
                controlnet_block = nn.Conv(output_channel, 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(controlnet_block)
        self.down_blocks = down_blocks
        self.controlnet_down_blocks = controlnet_down_blocks
        mid_block_channel = block_out_channels[-1]
        self.mid_block = FlaxUNetMidBlock2DCrossAttn(in_channels=mid_block_channel, dropout=self.dropout, num_attention_heads=num_attention_heads[-1], use_linear_projection=self.use_linear_projection, dtype=self.dtype)
        self.controlnet_mid_block = nn.Conv(mid_block_channel, kernel_size=(1, 1), padding='VALID', kernel_init=nn.initializers.zeros_init(), bias_init=nn.initializers.zeros_init(), dtype=self.dtype)
    def __call__(self, sample: jnp.ndarray, timesteps: Union[jnp.ndarray, float, int], encoder_hidden_states: jnp.ndarray, controlnet_cond: jnp.ndarray, conditioning_scale: float=1.0,
    return_dict: bool=True, train: bool=False) -> Union[FlaxControlNetOutput, Tuple[Tuple[jnp.ndarray, ...], jnp.ndarray]]:
        """Returns:"""
        channel_order = self.controlnet_conditioning_channel_order
        if channel_order == 'bgr': controlnet_cond = jnp.flip(controlnet_cond, axis=1)
        if not isinstance(timesteps, jnp.ndarray): timesteps = jnp.array([timesteps], dtype=jnp.int32)
        elif isinstance(timesteps, jnp.ndarray) and len(timesteps.shape) == 0:
            timesteps = timesteps.astype(dtype=jnp.float32)
            timesteps = jnp.expand_dims(timesteps, 0)
        t_emb = self.time_proj(timesteps)
        t_emb = self.time_embedding(t_emb)
        sample = jnp.transpose(sample, (0, 2, 3, 1))
        sample = self.conv_in(sample)
        controlnet_cond = jnp.transpose(controlnet_cond, (0, 2, 3, 1))
        controlnet_cond = self.controlnet_cond_embedding(controlnet_cond)
        sample += controlnet_cond
        down_block_res_samples = (sample,)
        for down_block in self.down_blocks:
            if isinstance(down_block, FlaxCrossAttnDownBlock2D): sample, res_samples = down_block(sample, t_emb, encoder_hidden_states, deterministic=not train)
            else: sample, res_samples = down_block(sample, t_emb, deterministic=not train)
            down_block_res_samples += res_samples
        sample = self.mid_block(sample, t_emb, encoder_hidden_states, deterministic=not train)
        controlnet_down_block_res_samples = ()
        for down_block_res_sample, controlnet_block in zip(down_block_res_samples, self.controlnet_down_blocks):
            down_block_res_sample = controlnet_block(down_block_res_sample)
            controlnet_down_block_res_samples += (down_block_res_sample,)
        down_block_res_samples = controlnet_down_block_res_samples
        mid_block_res_sample = self.controlnet_mid_block(sample)
        down_block_res_samples = [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=down_block_res_samples, mid_block_res_sample=mid_block_res_sample)
'''
    #################################################################################################################################################################################################
    # This code library is an adaptation of the original Transformers and was designed, developed and programmed by Sapiens Technology®.                                                            #
    # Any alteration and/or disclosure of this code without prior authorization is strictly prohibited and is subject to legal action that will be forwarded by the Sapiens Technology® legal team. #
    # This set of algorithms aims to download, train, fine-tune and/or infer large language models from various sources and slopes.                                                                 #
    #################################################################################################################################################################################################
'''
